Great, thanks so much! I added the train function in my original post!
I added the following to the train function but it doesn’t work.
  if step % 200 == 0 and not step == 0:
    valid_loss = evaluate() 
    print(f"Batch {step} / {len(train_data_loader)} | Valid Loss = {valid_loss}")
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), model_file)

Would be very happy if you could help me with this one, thanks!