Did you measure the individual gradient magnitudes for the loss w.r.t the network parameters? I don’t see it in the code how you meant the gradients are exploding in the validation phase. From your description it seems the training is fine i.e. loss is computed properly for training set and gradients backpropagating properly i.e. network is learning with loss reducing. But with validation your network generates nan output. By output you mean the loss that you are printing ? Or the actual output of your network which for example can be a softmax output?