fix test for pytorch 2.6

torchmd · Feb 20, 2025 · 06bdc7b · 06bdc7b
1 parent 767c68a
commit 06bdc7b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_priors.py b/tests/test_priors.py
@@ -186,7 +186,7 @@ def test_multiple_priors(dtype):
     with tempfile.NamedTemporaryFile() as f:
         torch.save(model, f)
         f.seek(0)
-        model2 = torch.load(f)
+        model2 = torch.load(f, weights_only=False)
         priors2 = model2.model.prior_model
         assert len(priors2) == 3
         assert isinstance(priors2[0], ZBL)