pytorch 常用准备阶段代码

原创于 2021-11-29 10:22:34 发布 · 1.2k 阅读

·

0

·

0 ·

CC 4.0 BY-SA版权

文章标签：

#python #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)