fixed test

CederGroupHub · Mar 1, 2024 · 2fb442d · 2fb442d
1 parent 3bc34b5
commit 2fb442d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
@@ -93,9 +93,9 @@ def test_trainer_composition_model(tmp_path: Path) -> None:
     new_chgnet = CHGNet.from_file(weights_path)
     for param in new_chgnet.composition_model.parameters():
         assert param.requires_grad is False
-    comparison = (
-        new_chgnet.composition_model.state_dict()["fc.weight"] == initial_weights
-    )
+    comparison = new_chgnet.composition_model.state_dict()["fc.weight"].to(
+        "cpu"
+    ) == initial_weights.to("cpu")
     expect = torch.ones_like(comparison)
     # Only Na and Cl should have updated
     expect[0][10] = 0