Print name in non-Array case in assertNestedAllClose(). (#164)

apple · Nov 4, 2023 · 32d439e · 32d439e
1 parent 8c13aa0
commit 32d439e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/axlearn/common/test_utils.py b/axlearn/common/test_utils.py
@@ -215,7 +215,7 @@ def assertNestedAllClose(self, a, b, atol=1e-6, rtol=1e-3):
                 self.assertEqual(a_value.shape, b_value.shape, msg=f"{a_name}")
                 assert_allclose(a_value, b_value, atol=atol, rtol=rtol, err_msg=f"{a_name}")
             else:
-                self.assertAlmostEqual(a_value, b_value)
+                self.assertAlmostEqual(a_value, b_value, msg=f"{a_name}")
 
     def assertNestedEqual(self, a, b):
         a_kv = flatten_items(a)