Hi there, im following tutorial of simple neural network with 2 hidden layers and I wanted to add MSE and CE charts at the end of file, however i have no idea what should i do, any ideas?
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as f
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import sklearn 
import numpy as np
import pandas as pd
import sklearn.metrics
import matplotlib.pyplot as plt

torch.manual_seed(101)

Transform = transforms.ToTensor()

train = torchvision.datasets.MNIST('', train=True, download=True, transform=Transform)
test = torchvision.datasets.MNIST('', train=False, download=True, transform=Transform)

trainset = DataLoader(train, batch_size=100, shuffle=True)
testset = DataLoader(test, batch_size=500, shuffle=False)

class Model(nn.Module):
    def __init__(self, input_size=784, output_size=10, layers=[120,84]):
        super().__init__()
        self.d1 = nn.Linear(input_size, layers[0])
        self.d2 = nn.Linear(layers[0], layers[1])
        self.d3 = nn.Linear(layers[1], output_size)
    
    def forward(self, X):
        X = f.relu(self.d1(X))
        X = f.relu(self.d2(X))
        X = self.d3(X)
        return f.log_softmax(X, dim=1)

model = Model()

ce = nn.CrossEntropyLoss()
mse = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

train_losses = []
train_mse_losses = []
test_losses = []
test_mse_losses = []

train_correct = []
test_correct = []

for i in range(10):
    trn_cor = 0
    tst_cor = 0
    for b, (X_train, y_train) in enumerate(trainset):
        b+=1
        y_pred = model(X_train.view(100, -1))
        loss = ce(y_pred, y_train)
        loss_mse = mse(y_pred, y_train)

        predicted = torch.max(y_pred.data, 1)[1]
        batch_cor = (predicted==y_train).sum()
        trn_cor += batch_cor

        optimizer.zero_grad()
        loss.backward()
        loss_mse.backward()
        optimizer.step()

        if b%600 == 0:
            print(f'epoch: {i:2} train loss: {loss.item():10.6f}')

    train_losses.append(loss)
    train_mse_losses.append(loss_mse)
    train_correct.append(trn_cor)

    with torch.no_grad():
        for b, (X_test, y_test) in enumerate(testset):
            y_val = model(X_test.view(500, -1))

            predicted = torch.max(y_val.data, 1)[1]
            tst_cor += (predicted == y_test).sum()
    
    loss = ce(y_val, y_test)
    loss_mse = mse(y_val, y_test)
    test_losses.append(loss)
    test_mse_losses.append(loss_mse)
    test_correct.append(tst_cor)

print(f'test acc: {test_correct[-1].item()*100/10000:.3f}%')

plt.subplot(3,1,1)
plt.plot(train_losses, label='training loss')
plt.plot(test_losses, label='validation loss')
plt.title('Loss at the end of each epoch')

plt.subplot(3,1,3)
plt.plot([t/600 for t in train_correct], label='training acc')
plt.plot([t/100 for t in train_correct], label='validation acc')
plt.title('Accuracy at the end of each epoch')

plt.legend()