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)
.
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)
.