Unable to run with cuequivariance acceleration

[dominicvarghese]

Hi,

I was trying to use the cuequivariance acceleration for MACE training. I followed the instructions from user guide and this is what my pip list looks like:

Package                       Version
----------------------------- ------------
ase                           3.26.0
ConfigArgParse                1.7.1
contourpy                     1.3.3
cuequivariance                0.7.0
cuequivariance-ops-cu12       0.7.0
cuequivariance-ops-torch-cu12 0.7.0
cuequivariance-torch          0.7.0
cycler                        0.12.1
e3nn                          0.4.4
filelock                      3.19.1
fonttools                     4.60.1
fsspec                        2025.9.0
gitdb                         4.0.12
GitPython                     3.1.45
h5py                          3.15.1
Jinja2                        3.1.6
kiwisolver                    1.4.9
lightning-utilities           0.15.2
lmdb                          1.7.5
mace-torch                    0.3.15
MarkupSafe                    2.1.5
matplotlib                    3.10.7
matscipy                      1.1.1
mpmath                        1.3.0
networkx                      3.5
numpy                         1.26.4
nvidia-cublas-cu12            12.6.4.1
nvidia-cuda-cupti-cu12        12.6.80
nvidia-cuda-nvrtc-cu12        12.6.77
nvidia-cuda-runtime-cu12      12.6.77
nvidia-cudnn-cu12             9.10.2.21
nvidia-cufft-cu12             11.3.0.4
nvidia-cufile-cu12            1.11.1.6
nvidia-curand-cu12            10.3.7.77
nvidia-cusolver-cu12          11.7.1.2
nvidia-cusparse-cu12          12.5.4.2
nvidia-cusparselt-cu12        0.7.1
nvidia-ml-py                  13.580.82
nvidia-nccl-cu12              2.27.5
nvidia-nvjitlink-cu12         12.6.85
nvidia-nvshmem-cu12           3.3.20
nvidia-nvtx-cu12              12.6.77
opt_einsum                    3.4.0
opt-einsum-fx                 0.1.4
orjson                        3.11.4
packaging                     25.0
pandas                        2.3.3
pillow                        11.3.0
pip                           25.3
platformdirs                  4.5.0
prettytable                   3.17.0
pyparsing                     3.2.5
python-dateutil               2.9.0.post0
python_hostlist               2.3.0
pytz                          2025.2
PyYAML                        6.0.3
scipy                         1.16.3
setuptools                    80.9.0
six                           1.17.0
smmap                         5.0.2
sympy                         1.14.0
torch                         2.9.1+cu126
torch-ema                     0.3
torchmetrics                  1.8.2
torchvision                   0.24.1+cu126
tqdm                          4.67.1
triton                        3.5.1
typing_extensions             4.15.0
tzdata                        2025.2
wcwidth                       0.2.14
wheel                         0.45.1

When I run the following on a v100 GPU node,

python -m mace.cli.run_train \
    --name="finetuned_MATPES-run" \
    --train_file="full_dataset.xyz" \
    --foundation_model="MACE-matpes-r2scan-omat-ft.model" \
    --pt_train_file="selected_1000_replay.xyz" \
    --multiheads_finetuning=True \
    --atomic_numbers="[8, 56]" \
    --energy_weight=1.0 \
    --compute_forces=True \
    --forces_weight=10.0 \
    --compute_stress=True \
    --stress_weight=5.0 \
    --E0s="e0s.json" \
    --energy_key="energy" \
    --forces_key="forces" \
    --stress_key="stress" \
    --max_num_epochs=18 \
    --config_type_weights '{"Default": 10.0}' \
    --weight_pt_head=1.0 \
    --lr=0.001 \
    --force_mh_ft_lr=True \
    --scheduler="ReduceLROnPlateau" \
    --lr_factor="0.5" \
    --scheduler_patience=10 \
    --scaling="rms_forces_scaling" \
    --batch_size=2 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --device=cuda \
    --enable_cueq=True \
    --valid_fraction=0.05 \
    --swa \
    --start_swa=12 \
    --swa_energy_weight=100 \
    --swa_forces_weight=10 \
    --swa_stress_weight=1 \
    --default_dtype="float64" \
    --seed=56 \
    --log_level="INFO" \
    --restart_latest \
    --error_table=PerAtomRMSE \
    --model="MACE"

I get the following error message:

/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/e3nn/o3/_wigner.py:10: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  _Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/cli/run_train.py:168: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  model_foundation = torch.load(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/modules/models.py:85: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
  "atomic_numbers", torch.tensor(atomic_numbers, dtype=torch.int64)
/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/tools/checkpoint.py:187: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  torch.load(f=checkpoint_info.path, map_location=device),
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/cli/run_train.py", line 1091, in <module>
    main()
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/cli/run_train.py", line 82, in main
    run(args)
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/cli/run_train.py", line 855, in run
    tools.train(
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/tools/train.py", line 195, in train
    valid_loss_head, eval_metrics = evaluate(
                                    ^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/tools/train.py", line 558, in evaluate
    output = model(
             ^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/modules/models.py", line 548, in forward
    node_feats, sc = interaction(
                     ^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/mace/modules/blocks.py", line 847, in forward
    mji = self.conv_tp(
          ^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_torch/operations/tp_channel_wise.py", line 251, in forward
    output = self.f(
             ^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_torch/primitives/segmented_polynomial.py", line 283, in forward
    return self.m(inputs, input_indices, output_shapes, output_indices)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_torch/primitives/segmented_polynomial_uniform_1d.py", line 254, in forward
    return tensor_product_uniform_1d_jit(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_torch/primitives/segmented_polynomial_uniform_1d.py", line 62, in tensor_product_uniform_1d_jit
    return torch.ops.cuequivariance.tensor_product_uniform_1d_jit(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_ops.py", line 1255, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_library/autograd.py", line 109, in autograd_impl
    result = Generated.apply(*args, Metadata(keyset, keyword_only_args))  # type: ignore[attr-defined]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_library/autograd.py", line 206, in new_apply
    result = orig_apply(*flat_args, metadata)  # type: ignore[misc]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/autograd/function.py", line 581, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_library/autograd.py", line 144, in new_forward
    result = orig_forward(ctx, *args)
             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_library/autograd.py", line 50, in forward
    result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_ops.py", line 848, in redispatch
    return self._handle.redispatch_boxed(keyset, *args, **kwargs)  # type: ignore[return-value]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_library/custom_ops.py", line 343, in backend_impl
    result = self._backend_fns[device_type](*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_compile.py", line 53, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 1044, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/_library/custom_ops.py", line 376, in wrapped_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops_torch/tensor_product_uniform_1d_jit.py", line 179, in _
    jit.run(
RuntimeError: CUDA error encountered at: file=/home/jenkins/agent/workspace/build-lib/cuequivariance_ops/cuda/equivariance/kernel.cuh line=413: call='cudaGetLastError()', Reason=cudaErrorNoKernelImageForDevice:no kernel image is available for execution on the device
Obtained 63 stack frames
#1 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x2a969) [0x151a5fedd969]
#2 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x5dc9b) [0x151a5ff10c9b]
#3 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x5e140) [0x151a5ff11140]
#4 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x4ee25) [0x151a5ff01e25]
#5 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x50bb1) [0x151a5ff03bb1]
#6 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops/lib/libcue_ops.so: kernelcatcher::equivariance::tensor_product_uniform_1d_jit::run_tensor_product_uniform_1d_jit(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, kernelcatcher::utils::Datatype, int, int, int, int, std::vector<kernelcatcher::equivariance::tensor_product_uniform_1d_jit::Dimension, std::allocator<kernelcatcher::equivariance::tensor_product_uniform_1d_jit::Dimension> > const&, std::vector<int, std::allocator<int> > const&, std::vector<std::vector<kernelcatcher::utils::BatchDimension, std::allocator<kernelcatcher::utils::BatchDimension> >, std::allocator<std::vector<kernelcatcher::utils::BatchDimension, std::allocator<kernelcatcher::utils::BatchDimension> > > > const&, std::vector<std::vector<int, std::allocator<int> >, std::allocator<std::vector<int, std::allocator<int> > > > const&, std::vector<int, std::allocator<int> > const&, std::vector<kernelcatcher::utils::Datatype, std::allocator<kernelcatcher::utils::Datatype> > const&, std::vector<std::vector<int, std::allocator<int> >, std::allocator<std::vector<int, std::allocator<int> > > > const&, std::vector<int, std::allocator<int> > const&, std::vector<int, std::allocator<int> > const&, std::vector<int, std::allocator<int> > const&, std::vector<int, std::allocator<int> > const&, std::vector<double, std::allocator<double> > const&, std::vector<int, std::allocator<int> > const&, std::vector<void*, std::allocator<void*> > const&, std::vector<unsigned long, std::allocator<unsigned long> > const&, bool, void*) +0x1006 [0x151a5ff077b6]
#7 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops_torch/_ext/cuequivariance_ops_torch_ext.cpython-311-x86_64-linux-gnu.so(+0x1b72d) [0x151a5fa2b72d]
#8 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops_torch/_ext/cuequivariance_ops_torch_ext.cpython-311-x86_64-linux-gnu.so(+0x1cff6) [0x151a5fa2cff6]
#9 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/cuequivariance_ops_torch/_ext/cuequivariance_ops_torch_ext.cpython-311-x86_64-linux-gnu.so(+0x226f3) [0x151a5fa326f3]
#10 in python: PyObject_Vectorcall +0x31 [0x51e5f1]
#11 in python: _PyEval_EvalFrameDefault +0x6a6 [0x511616]
#12 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#13 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#14 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#15 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#16 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#17 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#18 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#19 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#20 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#21 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x98317a) [0x151b76d8e17a]
#22 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0xcb7a5b) [0x151b770c2a5b]
#23 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0xcae634) [0x151b770b9634]
#24 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0xcaeabb) [0x151b770b9abb]
#25 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x3cc7ad) [0x151b767d77ad]
#26 in python() [0x528657]
#27 in python: _PyObject_MakeTpCall +0x26c [0x50400c]
#28 in python() [0x557b19]
#29 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#30 in python() [0x557928]
#31 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#32 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#33 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#34 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#35 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x7affb2) [0x151b76bbafb2]
#36 in python() [0x528680]
#37 in python: PyObject_Call +0xbd [0x542d3d]
#38 in python: _PyEval_EvalFrameDefault +0x9225 [0x51a195]
#39 in python() [0x557928]
#40 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#41 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#42 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#43 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#44 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0xcb70db) [0x151b770c20db]
#45 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0xcb7a01) [0x151b770c2a01]
#46 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(+0x5ce9594) [0x151b67759594]
#47 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0xa39175) [0x151b76e44175]
#48 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0xa394c0) [0x151b76e444c0]
#49 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so: torch::jit::_get_operation_for_overload_or_packet(std::vector<std::shared_ptr<torch::jit::Operator>, std::allocator<std::shared_ptr<torch::jit::Operator> > > const&, c10::Symbol, pybind11::args const&, pybind11::kwargs const&, bool, std::optional<c10::DispatchKey>) +0x38 [0x151b76e44528]
#50 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x92eb3c) [0x151b76d39b3c]
#51 in /jet/home/dvarghes/.conda/envs/mace-cueq/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x3cc7ad) [0x151b767d77ad]
#52 in python() [0x528657]
#53 in python: PyObject_Call +0xbd [0x542d3d]
#54 in python: _PyEval_EvalFrameDefault +0x9225 [0x51a195]
#55 in python: _PyFunction_Vectorcall +0x173 [0x538d03]
#56 in python: _PyObject_FastCallDictTstate +0x65 [0x5088d5]
#57 in python: _PyObject_Call_Prepend +0xe0 [0x540840]
#58 in python() [0x611ca7]
#59 in python: _PyObject_MakeTpCall +0x26c [0x50400c]
#60 in python: _PyEval_EvalFrameDefault +0x6a6 [0x511616]
#61 in python() [0x557928]
#62 in python: _PyEval_EvalFrameDefault +0x484b [0x5157bb]
#63 in python() [0x557928]

These were the modules loaded:

 module list

Currently Loaded Modules:
  1) allocations/1.0            3) intel-tbb/2021.10.0          5) intel-mkl/2023.2.0     7) openmpi/5.0.8-gcc13.3.1   9) cudnn/8.0.4
  2) psc.allocations.user/1.0   4) intel-compiler-rt/2023.2.1   6) gcc/13.3.1-p20240614   8) cuda/12.6.1              10) anaconda3/2024.10-1

So is this a CUDA version issue with the cuequivariance libraries?