from my experience nan occurs when the loss goes to infinity due to various reasons. print the loss(in your case cost) inside the data loader loop(where you printed 1 to 600) if the loss increases exponentially and reaches nan than decrease the learning rate and see again. normally if other things are right then decreasing learning rate will do. else there might be a problem with the model or with the data. hope it helps.