fix unit test

ewfuentes · Dec 17, 2024 · 78d9f8a · 78d9f8a
1 parent 0e2426e
commit 78d9f8a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/experimental/beacon_dist/utils_test.py b/experimental/beacon_dist/utils_test.py
@@ -123,7 +123,7 @@ def test_valid_configuration_loss_bad_prediction_yields_nonzero_loss(self):
         loss = valid_configuration_loss(batch.class_label, query, model_output)
 
         # Verification
-        self.assertGreater(loss, 0.0)
+        self.assertGreater(loss.item(), 0.0)
 
     def test_valid_configuration_loss_missing_exclusive_keypoints_are_valid(self):
         # Setup
@@ -136,7 +136,7 @@ def test_valid_configuration_loss_missing_exclusive_keypoints_are_valid(self):
         loss = valid_configuration_loss(batch.class_label, query, model_output)
 
         # Verification
-        self.assertAlmostEqual(loss, 0.0)
+        self.assertAlmostEqual(loss.item(), 0.0)
 
     def test_query_from_class_samples(self):
         # Setup