Solved by ptrblck in post #2 
              
              
                Could you try to isolate the iteration which causes the NaNs and check the input for invalid values as well as the loss in this and the last few iterations?