Clipping gradients to resolve NAN

[Ratish Puduppully]

Hi,
I have a network which uses LogSoftMax. It produces NANs and I intend to clip the gradients to resolve the issue.
A sample network:
output = module.forward(input)
loss = criterion.forward(output, target)
crit_grad = criterion.backward(output, target)
module_grad = module.backward(input, crit_grad)

Please confirm if I should clip the gradients of crit_grad or module_grad?

Regards,
Ratish