Solved by ptrblck in post #4 
              
              
                I think it’s easier to comment in the code: 
def train_model(model, criterion, optimizer, num_epochs):
    liveloss = PlotLosses()
    model = model.to(device)
    
    for epoch in range(num_epochs):
        logs = {}
        for phase in ['train', 'validation']:
            if phase == 'train':
  …