fix: UTs

anyangml · Feb 22, 2024 · e9d64cf · e9d64cf
1 parent 6da24c0
commit e9d64cf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py
@@ -254,7 +254,7 @@ def test_tf_consistent_with_ref(self):
         ret2 = self.extract_ret(ret2, self.RefBackend.TF)
         np.testing.assert_equal(data1, data2)
         for rr1, rr2 in zip(ret1, ret2):
-            np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol)
+            np.testing.assert_allclose(rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol)
 
     def test_tf_self_consistent(self):
         """Test whether TF is self consistent."""