Hotfix of the PET device

metatensor · Nov 26, 2024 · 6b330b4 · 6b330b4
1 parent 63a3ddb
commit 6b330b4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/metatrain/experimental/pet/trainer.py b/src/metatrain/experimental/pet/trainer.py
@@ -256,9 +256,9 @@ def train(
             else:
                 pet_model.hypers.TARGET_TYPE = "structural"
                 pet_model.TARGET_TYPE = "structural"
-            pet_model = pet_model.to(device=device, dtype=dtype)
         else:
             pet_model = PET(ARCHITECTURAL_HYPERS, 0.0, len(all_species))
+        pet_model = pet_model.to(device=device, dtype=dtype)
         num_params = sum([p.numel() for p in pet_model.parameters()])
         logging.info(f"Number of parameters: {num_params}")