pytorch实现mnist分类

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jun  9 19:40:53 2020

@author: 
"""

import torchvision
from matplotlib import pyplot as plt
import torch
from torchvision import datasets, transforms
from torch.autograd import Variable

#%%
transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5], std=[0.5])])
#transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5,0.5], std=[0.5,0.5])])
data_train=datasets.MNIST(root="/Users/didi/Documents/MNIST_master/",transform=transform,train=True,download=True)
data_test=datasets.MNIST(root="/Users/didi/Documents/MNIST_master/",transform=transform,train=False)
#%%
data_loader_train=torch.utils.data.DataLoader(dataset=data_train,batch_size=64,shuffle=True)
data_loader_test=torch.utils.data.DataLoader(dataset=data_test,batch_size=64,shuf