weird problem with pytorch lightning

[Holmes-Alan]

I tried a simple code of using EGNN for feature extraction, it works fine. But when I put it together in network training using pytorch lightning, I got the problem as following,

RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.

Can someone tell me where the problem is?