reset precision to double to fix tests

slaclab · Dec 20, 2024 · f34926f · f34926f
1 parent 6bb2e07
commit f34926f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/models/test_torch_model.py b/tests/models/test_torch_model.py
@@ -81,6 +81,8 @@ def test_precision(self, california_model):
         assert california_model.dtype == torch.double
         california_model.precision = "single"
         assert california_model.dtype == torch.float
+        # set back to double
+        california_model.precision = "double"
 
     def test_model_evaluate_single_sample(self, california_test_input_dict: dict, california_model):
         results = california_model.evaluate(california_test_input_dict)