half precision is super finicky during training, so I’m not surprised.
One thing I recommend trying is to do the forward + backward in half precision, but the optimizer step in float precision.
To do this, you might have to clone your parameters, and cast them to float32 and once forward+backward is over, you copy over the param .data and .grad into this float32 copy (and call optimizer.step on this float32 copy) and then copy back…
Other than that, I dont have a good idea of why adam + half is giving NaNs.