I have an NN architecture for molecular property prediction. The feature vectors for each atom go through its own MLP (based on atomic identity, i.e., atomic number). Then they are summed together to get the molecular energy. I use ModuleDict to keep track of those MLP. For molecular energies, I initialize an empty tensor with requires_grad = True, and then I populate the tensor going through each atom. My question is will it cause problems during backpropagation as pytorch doesn’t allow in-place operation (I use .data to bypass the problem but am not sure about the caveat of such bypassing). Below is a code snippet of the part that I am confused about,
atomic_energies = torch.zeros((self.batch_size, max_n_atoms),
            device=self.device, dtype=torch.float, requires_grad=True)
        
 for i in range(self.batch_size):
       for j in range(n_atoms):
            atomic_energies.data[i][j] = self.atomic_mlp_dict[species[i][j].item()](feature_matrix[i][j])

How should I verify that the gradients are being propagated correctly? Currently, the loss doesn’t change significantly for 10-20 epochs. How to debug this kind of problem in Pytorch? Any help is highly appreciated.