ML-IAP compute_forces failure when using MACE-MH-1

[MicPellegrino]

Hi,

I am testing the new MACE-MH-1 model (https://huggingface.co/mace-foundations/mace-mh-1). My build:

• LAMMPS 22Jul2025 (I pinned the release that worked for me so far),
• mace-torch 0.3.14
• torch 2.8.0+cu126 (CUDA Version: 12.6)
• cuequivariance 0.7.0

Upon running, I get this very long (and apparently low-level) error:

# We tested all the heads
python3 mace/mace/cli/create_lammps_model.py mace-mh-1.model --format=mliap --dtype=float64 --head='omat_pbe'
lammps-22Jul2025/build-mliap/lmp -k on g 1 -sf kk -pk kokkos newton on neigh half -in run-mliap.lammps

...

/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [64,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [65,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [66,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [67,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [68,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [69,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [70,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [71,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [72,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [73,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [74,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [75,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [76,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [77,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [78,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [79,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [80,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [81,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [82,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [83,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [84,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [85,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [86,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [87,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [88,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [89,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [90,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [91,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [92,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [93,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [94,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [95,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [32,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [33,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [34,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [35,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [36,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [37,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [38,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [39,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [40,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [41,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [42,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [43,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [44,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [45,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [46,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [47,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [48,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [49,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [50,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [51,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [52,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [53,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [54,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [55,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [56,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [57,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [58,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [59,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [60,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [61,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [62,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [58,0,0], thread: [63,0,0] Assertion `ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"` failed.
Traceback (most recent call last):
  File "mliap_unified_couple_kokkos.pyx", line 488, in mliap_unified_couple_kokkos.compute_forces_python_kokkos
  File "mliap_unified_couple_kokkos.pyx", line 470, in mliap_unified_couple_kokkos.MLIAPUnifiedInterfaceKokkos.compute_forces
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/mace/calculators/lammps_mliap_mace.py", line 160, in compute_forces
    _, atom_energies, pair_forces = self.model(batch)
                                    ^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/mace/calculators/lammps_mliap_mace.py", line 84, in forward
    out = self.model(
          ^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/mace/modules/models.py", line 537, in forward
    node_feats, sc = interaction(
                     ^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/mace/modules/blocks.py", line 1275, in forward
    mji = self.conv_tp(
          ^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_torch/operations/tp_channel_wise.py", line 251, in forward
    output = self.f(
             ^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_torch/primitives/segmented_polynomial.py", line 283, in forward
    return self.m(inputs, input_indices, output_shapes, output_indices)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_torch/primitives/segmented_polynomial_uniform_1d.py", line 254, in forward
    return tensor_product_uniform_1d_jit(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/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 "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/_ops.py", line 1243, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/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 "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/_library/autograd.py", line 206, in new_apply
    result = orig_apply(*flat_args, metadata)  # type: ignore[misc]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/autograd/function.py", line 576, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/_library/autograd.py", line 144, in new_forward
    result = orig_forward(ctx, *args)
             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/_library/autograd.py", line 50, in forward
    result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/_ops.py", line 836, in redispatch
    return self._handle.redispatch_boxed(keyset, *args, **kwargs)  # type: ignore[return-value]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/_library/custom_ops.py", line 344, in backend_impl
    result = self._backend_fns[device_type](*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/_compile.py", line 53, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 929, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/_library/custom_ops.py", line 377, in wrapped_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_ops_torch/tensor_product_uniform_1d_jit.py", line 179, in _
    jit.run(
RuntimeError: file=/home/jenkins/agent/workspace/build-lib/cuequivariance_ops/cuda/equivariance/kernel.cuh line=376: CUDA driver error encountered
Obtained 63 stack frames
#1 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x2a969) [0x7f294b42a969]
#2 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x582ee) [0x7f294b4582ee]
#3 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x4be7f) [0x7f294b44be7f]
#4 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x4ed9b) [0x7f294b44ed9b]
#5 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_ops/lib/libcue_ops.so(+0x50bb1) [0x7f294b450bb1]
#6 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/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 [0x7f294b4547b6]
#7 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_ops_torch/_ext/cuequivariance_ops_torch_ext.cpython-312-x86_64-linux-gnu.so(+0x1b879) [0x7f29868d4879]
#8 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_ops_torch/_ext/cuequivariance_ops_torch_ext.cpython-312-x86_64-linux-gnu.so(+0x1d142) [0x7f29868d6142]
#9 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/cuequivariance_ops_torch/_ext/cuequivariance_ops_torch_ext.cpython-312-x86_64-linux-gnu.so(+0x22873) [0x7f29868db873]
#10 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: PyObject_Vectorcall +0x50 [0x7f2eeba1e1c0]
#11 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x111dff) [0x7f2eeb911dff]
#12 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0x85294a) [0x7f2a39a5294a]
#13 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0xb584ab) [0x7f2a39d584ab]
#14 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0xb4c49f) [0x7f2a39d4c49f]
#15 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0xb4c955) [0x7f2a39d4c955]
#16 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0x3767d2) [0x7f2a395767d2]
#17 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x228a02) [0x7f2eeba28a02]
#18 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_MakeTpCall +0x2e3 [0x7f2eeba09783]
#19 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x257148) [0x7f2eeba57148]
#20 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x112a35) [0x7f2eeb912a35]
#21 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x256efb) [0x7f2eeba56efb]
#22 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x112a35) [0x7f2eeb912a35]
#23 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0x6a8f89) [0x7f2a398a8f89]
#24 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x228a29) [0x7f2eeba28a29]
#25 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_Call +0xb5 [0x7f2eeba3a885]
#26 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x112a35) [0x7f2eeb912a35]
#27 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x256efb) [0x7f2eeba56efb]
#28 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x112a35) [0x7f2eeb912a35]
#29 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0xb57a28) [0x7f2a39d57a28]
#30 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0xb58451) [0x7f2a39d58451]
#31 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so(+0x5b4224d) [0x7f2a2a74224d]
#32 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so: torch::jit::invokeOperatorFromPython(std::vector<std::shared_ptr<torch::jit::Operator>, std::allocator<std::shared_ptr<torch::jit::Operator> > > const&, pybind11::args const&, pybind11::kwargs const&, std::optional<c10::DispatchKey>) +0x394 [0x7f2a39b0ba04]
#33 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/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>) +0x1a9 [0x7f2a39b0bdc9]
#34 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0x81567a) [0x7f2a39a1567a]
#35 in /home/michele/mambaforge/envs/benchmark-mace/lib/python3.12/site-packages/torch/lib/libtorch_python.so(+0x3767d2) [0x7f2a395767d2]
#36 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x228a02) [0x7f2eeba28a02]
#37 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_Call +0xb5 [0x7f2eeba3a885]
#38 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x112a35) [0x7f2eeb912a35]
#39 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_FastCallDictTstate +0x1f7 [0x7f2eeba0c1e7]
#40 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_Call_Prepend +0xe9 [0x7f2eeba37d39]
#41 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x30ab00) [0x7f2eebb0ab00]
#42 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_MakeTpCall +0x2e3 [0x7f2eeba09783]
#43 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x111dff) [0x7f2eeb911dff]
#44 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x256efb) [0x7f2eeba56efb]
#45 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x112a35) [0x7f2eeb912a35]
#46 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x256efb) [0x7f2eeba56efb]
#47 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x112a35) [0x7f2eeb912a35]
#48 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_FastCallDictTstate +0x1f7 [0x7f2eeba0c1e7]
#49 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_Call_Prepend +0x69 [0x7f2eeba37cb9]
#50 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x30ab00) [0x7f2eebb0ab00]
#51 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_MakeTpCall +0x2e3 [0x7f2eeba09783]
#52 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x111dff) [0x7f2eeb911dff]
#53 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x257915) [0x7f2eeba57915]
#54 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x256f7d) [0x7f2eeba56f7d]
#55 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_Call +0x12b [0x7f2eeba3a8fb]
#56 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x112a35) [0x7f2eeb912a35]
#57 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x257915) [0x7f2eeba57915]
#58 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x256f7d) [0x7f2eeba56f7d]
#59 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_Call +0x12b [0x7f2eeba3a8fb]
#60 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x112a35) [0x7f2eeb912a35]
#61 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_FastCallDictTstate +0x291 [0x7f2eeba0c281]
#62 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0: _PyObject_Call_Prepend +0x69 [0x7f2eeba37cb9]
#63 in /home/michele/mambaforge/envs/benchmark-mace/lib/libpython3.12.so.1.0(+0x30ab00) [0x7f2eebb0ab00]

ERROR: Running mliappy unified compute_forces failure. (src/KOKKOS/mliap_unified_kokkos.cpp:82)
Last input line: run 1000
cudaFreeHost(Kokkos::Impl::CudaInternal::constantMemHostStagingPerDevice [cuda_device]) error( cudaErrorAssert): device-side assert triggered /home/michele/benchmark-mace/lammps-22Jul2025/lib/kokkos/core/src/Cuda/Kokkos_Cuda_Instance.cpp:648
Backtrace:
[0x7f2eee94bc09] Kokkos::Impl::save_stacktrace()
[0x7f2eee928300] Kokkos::Impl::host_abort(char const*)
[0x7f2eee95210b] Kokkos::Impl::cuda_internal_error_abort(cudaError, char const*, char const*, int)
[0x7f2eee95542a] Kokkos::Cuda::impl_finalize()
[0x7f2eee929bca] Kokkos::Impl::ExecSpaceManager::finalize_spaces()
[0x7f2eee92a06c] Kokkos::finalize()
[0x7f2eed84ba5d] LAMMPS_NS::KokkosLMP::finalize()
[0x61f0f64abd1a] 
[0x7f2eec029d90] 
[0x7f2eec029e40] __libc_start_main
[0x61f0f64abff5] 
[denerg33012:183277] *** Process received signal ***
[denerg33012:183277] Signal: Aborted (6)
[denerg33012:183277] Signal code:  (-6)
[denerg33012:183277] [ 0] /lib/x86_64-linux-gnu/libc.so.6(+0x42520)[0x7f2eec042520]
[denerg33012:183277] [ 1] /lib/x86_64-linux-gnu/libc.so.6(pthread_kill+0x12c)[0x7f2eec0969fc]
[denerg33012:183277] [ 2] /lib/x86_64-linux-gnu/libc.so.6(raise+0x16)[0x7f2eec042476]
[denerg33012:183277] [ 3] /lib/x86_64-linux-gnu/libc.so.6(abort+0xd3)[0x7f2eec0287f3]
[denerg33012:183277] [ 4] /home/michele/benchmark-mace/lammps-22Jul2025/build-mliap/liblammps.so.0(_ZN6Kokkos4Impl10host_abortEPKc+0x4d)[0x7f2eee92830d]
[denerg33012:183277] [ 5] /home/michele/benchmark-mace/lammps-22Jul2025/build-mliap/liblammps.so.0(_ZN6Kokkos4Impl25cuda_internal_error_abortE9cudaErrorPKcS3_i+0xfb)[0x7f2eee95210b]
[denerg33012:183277] [ 6] /home/michele/benchmark-mace/lammps-22Jul2025/build-mliap/liblammps.so.0(_ZN6Kokkos4Cuda13impl_finalizeEv+0x38a)[0x7f2eee95542a]
[denerg33012:183277] [ 7] /home/michele/benchmark-mace/lammps-22Jul2025/build-mliap/liblammps.so.0(_ZN6Kokkos4Impl16ExecSpaceManager15finalize_spacesEv+0x2a)[0x7f2eee929bca]
[denerg33012:183277] [ 8] /home/michele/benchmark-mace/lammps-22Jul2025/build-mliap/liblammps.so.0(_ZN6Kokkos8finalizeEv+0x2c)[0x7f2eee92a06c]
[denerg33012:183277] [ 9] /home/michele/benchmark-mace/lammps-22Jul2025/build-mliap/liblammps.so.0(_ZN9LAMMPS_NS9KokkosLMP8finalizeEv+0x2d)[0x7f2eed84ba5d]
[denerg33012:183277] [10] lammps-22Jul2025/build-mliap/lmp(+0xbd1a)[0x61f0f64abd1a]
[denerg33012:183277] [11] /lib/x86_64-linux-gnu/libc.so.6(+0x29d90)[0x7f2eec029d90]
[denerg33012:183277] [12] /lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0x80)[0x7f2eec029e40]
[denerg33012:183277] [13] lammps-22Jul2025/build-mliap/lmp(+0xbff5)[0x61f0f64abff5]
[denerg33012:183277] *** End of error message ***
Aborted (core dumped)

It is worth noting that:

• The MH-0 model works just fine;
• I get the same error regardless of the model head;
• I get the same error regardless of the model precision (single vs. double);
• We tested the model with the same LAMMPS build and the same environment both on a workstation (1GPU) and on a HPC cluster (2GPUs), and got the same error.

[ilyes319]

Hey @MicPellegrino,

Could you update to mace 0.3.15, that should be working. Also try Pytorch 2.9 and the latest cueq.

[MicPellegrino]

By "latest cueq" you mean 0.9?
The GitHub page of foundation models reports ">=v0.3.14" for mace, maybe you would want to update that.

[ilyes319]

yes 0.9

[MicPellegrino]

It works, I am now using:

• cuequivariance 0.9.1
• mace-torch 0.3.15 (pull latest commit of the mace github for model creation: e8a684f7448c386f798f9e8f363ad98922969d4e)
• torch 2.9.0+cu126
• LAMMPS 11 February 2026 release