Joseph_Santarcangelo:

LOSS.append(loss)


If you want to store the whole computational graph then its okay to use LOSS.append(loss);  but if you are just looking to store the loss value then use LOSS.append(loss.item()).



 Joseph_Santarcangelo:

LOSS=[]

n=1;
for epoch in range(5):
    for x,y in trainloader:
        yhat=model(x)
        loss=criterion(yhat,y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        LOSS.append(loss)



Also, shouldn’t you be accumulating the loss (for each batch) within each epoch and then appending it to loss array like this (below):
LOSS=[]
...
for epoch in range(5):
    running_loss = 0     #accumulates loss of each batch
    for x,y in trainloader:
        # do something......
        loss.backward()
        optimizer.step()
        running_loss += loss.item()   # accumulating loss
    
    LOSS.append(running_loss)   # saving final loss for each epoch