A question about cuda acceleration

[quva-lab]

A question about cuda acceleration

[ejnnr]

The nn subpackage uses Pytorch internally for everything, so it supports CUDA acceleration. You can move GeometricTensors to GPU with GeometricTensor.to. EquivariantModules are just Pytorch Modules so they can also be moved to GPU with .to(device).

The other subpackages use numpy and are therefore not CUDA accelerated (also see
#39 (comment)). But that part of the code is only run once when you initialize your network. So all of the training and inference can be done on GPU.

[Gabri95]

Hi @kgavrilyuk

Thanks @ejnnr for the faster reply :)
Indeed, CUDA is supported in the nn subpackage, so once you have built your neural network, you should be able to use CUDA acceleration as you usually do in PyTorch

Best,
Gabriele