Hello,
Could you try the following code:
criterion = nn.NLLLoss()
optimizer = torch.optim.Adam(model.parameters(). lr=learning_rate, weight_decay=1e-5)
test_op = model(test_img.cuda())
criterion(test_op, test_lb.cuda())