Add None check before np.testing.assert_allclose

njzjz · Nov 28, 2024 · 77cf60d · 77cf60d
1 parent b0520b0
commit 77cf60d
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py
@@ -184,11 +184,12 @@ def test_forward(self) -> None:
 
         for kk in ret[0]:
             # ensure the first frame and the second frame are the same
-            np.testing.assert_allclose(
-                ret[0][kk][0],
-                ret[0][kk][1],
-                err_msg=f"compare {kk} between frame 0 and 1",
-            )
+            if ret[0][kk] is not None:
+                np.testing.assert_allclose(
+                    ret[0][kk][0],
+                    ret[0][kk][1],
+                    err_msg=f"compare {kk} between frame 0 and 1",
+                )
 
             subret = []
             for rr in ret: