add debugging statement

mlcommons · Nov 18, 2023 · b1c35d4 · b1c35d4
1 parent d7d4638
commit b1c35d4
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/modeldiffs/torch2jax_utils.py b/tests/modeldiffs/torch2jax_utils.py
@@ -106,6 +106,7 @@ def diff(self):
         if s_p == s_j:
           count += 1
         else:
+          print('Difference in pytorch and jax key:')
           print(k, s_p, s_j)
       print(f'Number of values with identical shapes: {count}')