pytorch 常用准备阶段代码

import numpy as np

from tqdm import tqdm
import random
import torch.optim as optim
import os
import torch
import torch.nn as nn
os.environ["CUDA_VISIBLE_DEVICES"]="0"
use_cuda = torch.cuda.is_available()
print('use_cuda', use_cuda)
device = torch.device("cuda:0" if use_cuda else "cpu")

np.random.seed(0)
torch.manual_seed(0)
random.seed(0)