So it seems that the torch.exp call is returning an invalid gradient. Could you print out the output of the torch.exp(-x) call where x is defined as ((input_unfold - weight_flat).sum(dim=2))? and possibly input_unfold - weight_flat?
I have a feeling torch.exp is return an infinity which when passed to the backward causes the NaN.
Did you check that you’re using torch.exp(-x) rather than torch.exp(x), the minus sign in the exponent is needed.



 Sandeep_Kumar_GITAM:

            output = (1 + ((input_unfold - weight_flat).sum(dim=2)).exp()).reciprocal()



This should be output = (1 + ((input_unfold - weight_flat).sum(dim=2)).mul(-1).exp()).reciprocal(), in order for it to be a sigmoid function!