Lau:

criterion = nn.CrossEntropyLoss()
optimizer= optim.Adam(final_model.parameters(), lr = 0.001)

if torch.cuda.is_available():
    final_model.cuda()  



Thank you very much @ptrblck!! As always, you’re super helpful! Have a great day!