Open qiujunhan opened 4 years ago
报错 内容
Traceback (most recent call last):
File ".\train.py", line 68, in
train.py中修改:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
cnn = CNN().to(device) images = Variable(images).to(device) labels = Variable(labels.float()).to(device)
![屏幕截图 2021-02-23 103023](https://user-images.githubusercontent.com/21073721/108795942-7b706200-75c2-11eb-8f88-68525acb7768.png
找了好久资料都不知道怎么用GPU训练,一些设置也改了
-- coding: UTF-8 --
import torch import torch.nn as nn from torch.autograd import Variable import datasets from models import * import torch_util import os, shutil import argparse import test import torchvision import settings
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
Hyper Parameters
num_epochs = 300 batch_size = 20 learning_rate = 0.001
device = torch_util.select_device() device = torch.device("cuda:0")
def main(args): cnn = CNN().to(device)