fixing test

bussilab · Nov 28, 2024 · c22fe83 · c22fe83
1 parent cdbb387
commit c22fe83
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/tests/test_loss_and_minimizer.py b/tests/test_loss_and_minimizer.py
@@ -10,6 +10,11 @@ def assertEqualObjs(self, obj1, obj2):
 
         print(obj1, obj2)
 
+        if isinstance(obj1, np.ndarray):
+            if obj1.shape == (1,): obj1 = obj1[0]
+        if isinstance(obj2, np.ndarray):
+            if obj2.shape == (1,): obj2 = obj2[0]
+
         if isinstance(obj1, dict) and isinstance(obj2, dict):
             self.assertSetEqual(set(obj1.keys()), set(obj2.keys()))
             for k in obj1.keys():

diff --git a/tests/test_mdrefinement.py b/tests/test_mdrefinement.py
@@ -9,6 +9,11 @@ def assertEqualObjs(self, obj1, obj2, tol = 1e-4, if_relative = False):
 
         print(obj1, obj2)
 
+        if isinstance(obj1, np.ndarray):
+            if obj1.shape == (1,): obj1 = obj1[0]
+        if isinstance(obj2, np.ndarray):
+            if obj2.shape == (1,): obj2 = obj2[0]
+
         if isinstance(obj1, dict) and isinstance(obj2, dict):
             self.assertSetEqual(set(obj1.keys()), set(obj2.keys()))
             for k in obj1.keys():