Thanks a lot for your reply Alban.
I checked the position variable, I only changed it in one place within the forward() func:
positions = positions.unsqueeze(-1).expand(-1, -1, fdim)

and the same error still exists. If I don’t use positions as an input argument, actually this error disappears.
Also, I’m wondering what’s version of a variable means. I got different version numbers if I change some of the operations on positions.
Thanks again.