有没有大佬会改成用GPU计算的

[晨昏线ci]

import numpy import torch import torch.nn as nn import torch.optim as optim import matplotlib.pyplot as plt import torchvision import torchvision.transforms as transforms import time time_open = time.time() transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean = [0.5],std = [0.5])]) #train_data = torchvision.datasets.MNIST(root = './data',download = True,transform = transforms.ToTensor(),train = True) #此函数直接用transforms.ToTensor() 像素值在0.5---1之间 不是正太分布 颜色是灰到白 #test_data = torchvision.datasets.MNIST(root = './data',download = True,transform = transforms.ToTensor(),train = True) train_data = torchvision.datasets.MNIST(root = './data',download = True,transform = transform,train = True) test_data = torchvision.datasets.MNIST(root = './data',download = True,transform = transform,train = True) #调用transform函数 将像素值调整到0--1之间的正太分布 黑为0百为1 灰为0.5 train_data_sets = torch.utils.data.DataLoader(dataset = train_data,shuffle = True,batch_size = 64) test_data_sets = torch.utils.data.DataLoader(dataset = test_data,shuffle = True,batch_size = 64) images,label = next(iter(train_data_sets)) img = torchvision.utils.make_grid(images) img = img.numpy().transpose(1,2,0) img = img/2+0.5 print([label[i].item() for i in range(64)]) plt.imshow(img) plt.show() class Net(nn.Module): def __init__(self): super(Net,self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(1,64,kernel_size = 3,stride = 1,padding = 1), nn.ReLU(), nn.Conv2d(64,128,kernel_size = 3,stride = 1,padding = 1), nn.ReLU(), nn.MaxPool2d(kernel_size = 2,stride = 2)) self.dense = nn.Sequential( nn.Linear(14*14*128,1024), nn.ReLU(), nn.Dropout(p=0.5), #nn.Linear(1024,512), #nn.ReLU(), nn.Linear(1024,10)) def forward(self,x): x = self.conv1(x) x = x.view(-1,14*14*128) x = self.de

[linxxx3]

import time 下面加一行

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")