16 bit precision has a tendency to truncate values to 0 more often than 32 bit. this might be triggering a inf which might lead to a downwind nan.  you can trigger a breakpoint by

checking the output of the network (i.e. end of forward pass) for inf or nan. and stopping if you encounter it
checking the first layer gradients for nan/inf ( i.e. end of the backward pass), and stopping the code (before optimizer.step is actually called and ends up corrupting your weights) if this happens

then you can inspect the conditions around which the nans arise ( was it in the forward or backward, had something to do with data/ float 16?)