I modify and combine  the code on the internet recently ,
Did it seem to be wrong about  train auccuracy , train loss, test auccuracy , test auccuracy?
or others improve suggestion?
thanks
code:
 import torch
 import torch.nn as nn
 import matplotlib.pyplot as plt
 from torch.utils.data import DataLoader
 from torchvision import datasets, transforms

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_file = datasets.MNIST(root='./data/',train=True,transform=transforms.ToTensor(),download=True)

test_file = datasets.MNIST(root='./dataset/',train=False,transform=transforms.ToTensor(),download=True)

train_loader = DataLoader(dataset=train_file,batch_size=128,shuffle=True)

test_loader = DataLoader(dataset=test_file,batch_size=128,shuffle=False)

class CNN(nn.Module):

    def __init__(self):

        super(CNN,self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Linear(64 * 7 * 7, 10)

    def forward(self,x):
        x = self.conv(x)
        x = x.view(x.size(0), -1) 
        y = self.fc(x)
        return y

model = CNN().to(device)
print(model)

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3,momentum=0.9)
loss_fn = nn.CrossEntropyLoss()

def train(dataloader , model ,loss_fn, optimizer):

    size = len(dataloader.dataset)

    num_batches = len(dataloader)

    train_loss = 0

    train_Acc = 0

    for data , targets in train_loader:

        data = data.to(device)

        targets = targets.to(device)

        # Compute prediciton error

        pred = model(data)

        loss = loss_fn(pred,targets)

        # Backpropagation

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

   

        train_loss += loss.item()

        train_Acc += (pred.argmax(1) == targets).type(torch.float).sum().item()

    train_loss /= num_batches

    train_Acc /= size

    print('Train Auccuracy: %.1f , Train loss: %.4f \n' %(100*train_Acc , train_loss) )

def test(dataloader, model ,loss_fn):

    size = len(dataloader.dataset)

    num_batches = len(dataloader)

    model.eval()

    test_loss = 0 

    test_acc = 0

    with torch.no_grad():

        for data, targets in test_loader:

            data,targets = data.to(device),targets.to(device)

            pred = model(data)

            test_loss += loss_fn(pred,targets).item()

            test_acc += (pred.argmax(1) == targets).type(torch.float).sum().item()

    test_loss /= num_batches

    test_acc /= size

    print('Test Auccuracy: %.1f , Test loss: %.4f \n' %(100*test_acc , test_loss) )

epochs = 5

for t in range(epochs):

    print(f"Epoch {t+1}\n-------------------------------")

    train(train_loader, model, loss_fn, optimizer)

    test(test_loader, model, loss_fn)

print("Done!")