[SE3Transformer] tests failed on GPU?

[xingruiyang]

Related to DGLPytorch/SE3Transformer

Describe the bug failed running tests/test_equivariance.py

Traceback (most recent call last):
  File "test.py", line 27, in <module>
    test_equivariance()
  File "DeepLearningExamples/DGLPyTorch/DrugDiscovery/SE3Transformer/tests/test_equivariance.py", line 81, in test_equivariance
    assert torch.allclose(out2['0'], out1['0'], atol=TOL), \
AssertionError: type-0 features should be invariant 0.15505647659301758

but the same code runs OK with only CPU

To Reproduce all three tests failed including test_equivariance, test_equivariance_pooled and def test_invariance_pooled

Environment Please provide at least:

• Pytorch version: torch==1.11.0+cu113, dgl-cu110==0.6.1
• GPUs in the system: RTX 3090
• CUDA version: 11.3
• CUDA driver version (e.g. 418.67): 470.103.01