So, to give a bit more of context, I think I can show a MWE. Essentially I’m running the RBM in https://github.com/odie2630463/Restricted-Boltzmann-Machines-in-pytorch 41, with just a fundamental modification. To compute free_energy one exponentiates wx_b, adds 1, and takes the log. If wx_b is too large, the exponentiation gives inf, which after adding 1 and taking the log remains as inf, while the result should be negligibly close to wx_b. To avoid this, I modified free_energy by
 def free_energy(self,v):
        vbias_term = v.mv(self.v_bias)
        wx_b = F.linear(v,self.W,self.h_bias)
        hidden_term = wx_b.exp().add(1).log()
        if (hidden_term == np.inf).sum().data.cpu().numpy() != 0:
             hidden_term[hidden_term == np.inf] = wx_b[hidden_term == np.inf]
        hidden_term = hidden_term.sum(1)
        return (-hidden_term - vbias_term).mean()

Also, a minor modification was to substitute the optimizer by Adam and set the learning rate to 5e-2. Anyway, in my more complicated model SGD is giving the NaN in the gradients as well. Also I modified the number of units in the hidden layer to 50.
Even with this fix, I see a NaN appearing in the gradients at some specific point in the training. However, neither the input values, nor the intermediate calculations, nor the value of the loss function the gradients are calculated from seem to have problems in the step just before of getting the NaNgradient.
If it may be useful, I have observed that the first time I get a NaN, it simultaneously appears in just one gradient for the biases of the hidden units (i.e., in one cell of params[2].grad), and one full column for the weights (in params[0].grad), actually the column indexed by the same index as the NaN hidden unit bias gradient. This is, I’m getting NaNs simultaneously in the gradients dL/d(hbias_i) and dL/d(W_{i j}) for just one specific i.
Any suggestion will be much appreciated, this is already driving me crazy…