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()
```