diff --git a/torch_sim/state.py b/torch_sim/state.py
index ccaffb27..cf49a310 100644
--- a/torch_sim/state.py
+++ b/torch_sim/state.py
@@ -172,11 +172,7 @@ def n_atoms(self) -> int:
     @property
     def n_atoms_per_system(self) -> torch.Tensor:
         """Number of atoms per system."""
-        return (
-            self.system_idx.bincount()
-            if self.system_idx is not None
-            else torch.tensor([self.n_atoms], device=self.device)
-        )
+        return self.system_idx.bincount()
 
     @property
     def n_systems(self) -> int: