From fdb7ebeea0eb26c9adf20641f53fa5938a9a07a5 Mon Sep 17 00:00:00 2001 From: cw-tan Date: Wed, 19 Nov 2025 20:39:39 -0500 Subject: [PATCH] fix PBC on GPU bug in `state_to_atoms` --- torch_sim/io.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torch_sim/io.py b/torch_sim/io.py index 253b29ed..eee9e480 100644 --- a/torch_sim/io.py +++ b/torch_sim/io.py @@ -53,6 +53,7 @@ def state_to_atoms(state: "ts.SimState") -> list["Atoms"]: cell = state.cell.detach().cpu().numpy() # Shape: (n_systems, 3, 3) atomic_numbers = state.atomic_numbers.detach().cpu().numpy() system_indices = state.system_idx.detach().cpu().numpy() + pbc = state.pbc.detach().cpu().numpy() atoms_list = [] for sys_idx in np.unique(system_indices): @@ -65,7 +66,7 @@ def state_to_atoms(state: "ts.SimState") -> list["Atoms"]: symbols = [chemical_symbols[z] for z in system_numbers] atoms = Atoms( - symbols=symbols, positions=system_positions, cell=system_cell, pbc=state.pbc + symbols=symbols, positions=system_positions, cell=system_cell, pbc=pbc ) atoms_list.append(atoms)