Getting Nan after first iteration with custom loss

4

Could you check the gradients in the layers which have the NANs after the update?
You can print them with print(model.fc1.weight.grad).