From 32d439e99335f0d5d1cab7afb5574319eaf4abd1 Mon Sep 17 00:00:00 2001
From: apghml <143655008+apghml@users.noreply.github.com>
Date: Sat, 4 Nov 2023 04:12:14 -0700
Subject: [PATCH] Print name in non-Array case in `assertNestedAllClose()`.
 (#164)

---
 axlearn/common/test_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/axlearn/common/test_utils.py b/axlearn/common/test_utils.py
index 6139de82..15298507 100644
--- 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)