Hi ptrblck,
Thanks for your reply!
I think maybe one reason is that pytorch cannot track the gradient process once I use a function in the “for loop”? For example, I use RigidityExp function in the for loop as below:
for i in range(X.shape[0]):
            output[i,:] = RigidityExp.apply(X[i,:].reshape(1,-1).double(), self.eta, self.beta, self.cutoff)

I’m new to pytorch so that’s just my guess. My datasets do not have a decent shape so I will try to custom my dataset to see if it works.
Anyway, I really appreciate your help! And also I will let you know if custom dataset works.
Best,
qflm