How to Implement a Custom Loss Function with Keras for a Sparse Dataset
From the mentioned problems you are facing, this seems like a problem of exploding gradients. The exploding gradients problem can be identified by:


The model is unable to get traction on your training data (e.g. poor loss).
The model is unstable, resulting in large changes in loss from update to update.
The model loss goes to NaN during training.


More about Exploding gradient problem can be found at this article

I would suggest you to use some gradient clipping technique in you code and this will remove the NaN generation during model training.