Missing device in `FixCom` / `get_centers_of_mass`

[sklenard]

Some torch.zeros() tensors are not placed on the correct device when adding a FixCom constraint or when calling get_centers_of_mass() leading to a RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! when running on GPU. This is somehow related to issue #371.

Minimal reproduction:

import torch
from ase.build import bulk

import torch_sim as ts
from torch_sim.constraints import FixCom
from torch_sim.models.lennard_jones import LennardJonesModel
from torch_sim.transforms import get_centers_of_mass
from torch_sim.units import MetalUnits

n_steps = 2000
dt = 0.001            # ps
temperature = 1000.0  # K
dtype = torch.float64
device = "cuda"

atoms = bulk("Ar", "fcc", a=5.26, cubic=True) * (2, 2, 2)  # 32 atoms
atoms.rattle(stdev=0.05, seed=0)

calc = LennardJonesModel(
    sigma=3.405,
    epsilon=0.0104,
    dtype=dtype,
    device=device,
    compute_forces=True,
)

# --- 3. Build SimState from atoms, attach FixCom, record initial COM ---
state = ts.initialize_state(atoms, device=device, dtype=dtype)
state.constraints = [FixCom([0])]

initial_com = get_centers_of_mass(
    state.positions,
    state.masses,
    state.system_idx,
    state.n_systems,
)

state = ts.integrate(
    system=state,
    model=calc,
    n_steps=n_steps,
    timestep=dt,
    temperature=temperature,
    integrator=ts.Integrator.nvt_nose_hoover,
    pbar=True,
)

final_com = get_centers_of_mass(
    state.positions,
    state.masses,
    state.system_idx,
    state.n_systems,
)

print("Initial COM:", initial_com.tolist())
print("Final   COM:", final_com.tolist())
print("COM drift  :", (final_com - initial_com).abs().max().item())

[CompRhys]

Thanks for the report! would you be willing to submit a PR to fix this?

[sklenard]

Sure !