Sandeep_Kumar_GITAM:

((input_unfold - weight_flat).sum(dim=2))


whatever’s computed within this is causing torch.exp() to return a NaN value. I see you’re using sigmoid as your given function? If so, you’re missing the minus sign on your exponent.
Sigmoid is defined as 1/(1 + torch.exp(-x)) whereas you have 1/(1 + exp(x)) at the moment where x is defined as ((input_unfold - weight_flat).sum(dim=2)).
You could just pass it into torch.sigmoid directly, via,
output = torch.sigmoid( ((input_unfold - weight_flat).sum(dim=2)) )