fix device om calculator init

src/torchpme/_utils.py

@@ -11,26 +11,20 @@ def _validate_parameters(
     neighbor_distances: torch.Tensor,
     smearing: Union[float, None],
     dtype: torch.dtype,
-    device: Union[str, torch.device],
+    device: torch.device,
 ) -> None:
     if positions.dtype != dtype:
         raise TypeError(
             f"type of `positions` ({positions.dtype}) must be same as the class "
             f"type ({dtype})"
         )
 
-    if isinstance(device, torch.device):
-        device = device.type
-
-    if positions.device.type != device:
+    if positions.device != device:
         raise ValueError(
             f"device of `positions` ({positions.device}) must be same as the class "
             f"device ({device})"
         )
 
-    # We use `positions.device` because it includes the device type AND index, which the
-    # `device` parameter may lack
-
     # check shape, dtype and device of positions
     num_atoms = len(positions)
     if list(positions.shape) != [len(positions), 3]:
@@ -51,7 +45,7 @@ def _validate_parameters(
             f"type of `cell` ({cell.dtype}) must be same as the class ({dtype})"
         )
 
-    if cell.device != positions.device:
+    if cell.device != device:
         raise ValueError(
             f"device of `cell` ({cell.device}) must be same as the class ({device})"
         )
@@ -85,7 +79,7 @@ def _validate_parameters(
             f"type of `charges` ({charges.dtype}) must be same as the class ({dtype})"
         )
 
-    if charges.device != positions.device:
+    if charges.device != device:
         raise ValueError(
             f"device of `charges` ({charges.device}) must be same as the class "
             f"({device})"
@@ -99,7 +93,7 @@ def _validate_parameters(
             "structure"
         )
 
-    if neighbor_indices.device != positions.device:
+    if neighbor_indices.device != device:
         raise ValueError(
             f"device of `neighbor_indices` ({neighbor_indices.device}) must be "
             f"same as the class ({device})"
@@ -112,7 +106,7 @@ def _validate_parameters(
             f"{list(neighbor_indices.shape)} and {list(neighbor_distances.shape)}"
         )
 
-    if neighbor_distances.device != positions.device:
+    if neighbor_distances.device != device:
         raise ValueError(
             f"device of `neighbor_distances` ({neighbor_distances.device}) must be "
             f"same as the class ({device})"

src/torchpme/calculators/calculator.py

@@ -46,7 +46,7 @@ def __init__(
                 f"Potential must be an instance of Potential, got {type(potential)}"
             )
 
-        self.device = torch.get_default_device() if device is None else device
+        self.device = torch.get_default_device() if device is None else torch.device(device)
         self.dtype = torch.get_default_dtype() if dtype is None else dtype
 
         if self.dtype != potential.dtype:

src/torchpme/potentials/potential.py

@@ -43,7 +43,7 @@ def __init__(
     ):
         super().__init__()
         self.dtype = torch.get_default_dtype() if dtype is None else dtype
-        self.device = torch.get_default_device() if device is None else device
+        self.device = torch.get_default_device() if device is None else torch.device(device)
         if smearing is not None:
             self.register_buffer(
                 "smearing", torch.tensor(smearing, device=self.device, dtype=self.dtype)

src/torchpme/tuning/tuner.py

@@ -91,7 +91,7 @@ def __init__(
                 f"Only exponent = 1 is supported but got {exponent}."
             )
 
-        self.device = torch.get_default_device() if device is None else device
+        self.device = torch.get_default_device() if device is None else torch.device(device)
         self.dtype = torch.get_default_dtype() if dtype is None else dtype
 
         _validate_parameters(
@@ -295,7 +295,7 @@ def __init__(
         super().__init__()
 
         self.dtype = torch.get_default_dtype() if dtype is None else dtype
-        self.device = torch.get_default_device() if device is None else device
+        self.device = torch.get_default_device() if device is None else torch.device(device)
 
         _validate_parameters(
             charges=charges,

tests/calculators/test_workflow.py

@@ -60,7 +60,7 @@
 class TestWorkflow:
     def cscl_system(self, device=None, dtype=None):
         """CsCl crystal. Same as in the madelung test"""
-        device = torch.get_default_device() if device is None else device
+        device = torch.get_default_device() if device is None else torch.device(device)
         dtype = torch.get_default_dtype() if dtype is None else dtype
 
         positions = torch.tensor(

tests/helpers.py

@@ -15,7 +15,7 @@
 
 
 def define_crystal(crystal_name="CsCl", dtype=None, device=None):
-    device = torch.get_default_device() if device is None else device
+    device = torch.get_default_device() if device is None else torch.device(device)
     dtype = torch.get_default_dtype() if dtype is None else dtype
 
     # Define all relevant parameters (atom positions, charges, cell) of the reference

tests/metatensor/test_workflow_metatensor.py

@@ -54,7 +54,7 @@
 )
 class TestWorkflow:
     def system(self, device=None, dtype=None):
-        device = torch.get_default_device() if device is None else device
+        device = torch.get_default_device() if device is None else torch.device(device)
         dtype = torch.get_default_dtype() if dtype is None else dtype
 
         system = mts_atomistic.System(

tests/tuning/test_tuning.py

@@ -22,7 +22,7 @@
 
 
 def system(device=None, dtype=None):
-    device = torch.get_default_device() if device is None else device
+    device = torch.get_default_device() if device is None else torch.device(device)
     dtype = torch.get_default_dtype() if dtype is None else dtype
 
     charges = torch.ones((4, 1), dtype=dtype, device=device)