Description

In line 17 of e3nn/math/_normalize_activation.py, I changed from

    z = torch.randn(1_000_000, generator=gen, dtype=torch.float64, device="cpu").to(dtype=dtype, device=device)

Motivation and Context

Background: I'm using E3NN in a training code written with pytorch-lightning. During the training process, I encountered the following error:

  File "/opt/conda/lib/python3.10/site-packages/e3nn/nn/_fc.py", line 73, in __init__
    act = normalize2mom(act)
  File "/opt/conda/lib/python3.10/site-packages/e3nn/math/_normalize_activation.py", line 43, in __init__
    cst = moment(f, 2, dtype=torch.float64, device="cpu").pow(-0.5).item()
  File "/opt/conda/lib/python3.10/site-packages/e3nn/math/_normalize_activation.py", line 17, in moment
    z = torch.randn(1_000_000, generator=gen, dtype=torch.float64).to(dtype=dtype, device=device)
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/_device.py", line 77, in __torch_function__
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/fabric/utilities/init.py", line 53, in __torch_function__
    return func(*args, **kwargs)
RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'

From the code I assume the process is to create a random number array on CPU and then copy to GPU. Under the environment of the pytorch-lightning, it must be trying the create the random number array on GPU and failed since everything else has been placed explicitly on CPU.
Therefore I propose to add the device="cpu" to the torch.randn function call, to avoid this error.

How Has This Been Tested?

I have tested the fixed code training my own model. It successfully fixed the error.
Since the edit is minimal, it should be safe from introducing bugs.

Checklist:

• I have read the CONTRIBUTING document.
• My code follows the code style of this project.
• I have updated the documentation (if relevant).
• I have added tests that cover my changes (if relevant).
• The modified code is cuda compatible (github tests don't test cuda) (if relevant).
• I have updated the CHANGELOG.