Fix exception handling in torch._dynamo.utils.same and add corresponding test

test/dynamo/test_same_allclose.py

@@ -0,0 +1,25 @@
+import unittest
+import torch
+from torch._dynamo.testing import same
+
+class TestSameAllclose(unittest.TestCase):
+ def test_same_allclose_exception(self):
+ ref = torch.tensor([1.0, 2.0, 3.0])
+ res = torch.tensor([1.0, 2.0, float('nan')])
+
+ try:
+ same(ref, res, equal_nan=False)
+ except RuntimeError as e:
+ self.assertIn("An unexpected error occurred while comparing tensors with torch.allclose", str(e))
+ self.assertIn("RuntimeError: The size of tensor a (3) must match the size of tensor b (3) at non-singleton dimension 0", str(e))
+ else:
+ self.fail("RuntimeError not raised")
+
+ def test_same_allclose_no_exception(self):
+ ref = torch.tensor([1.0, 2.0, 3.0])
+ res = torch.tensor([1.0, 2.0, 3.0])
+
+ self.assertTrue(same(ref, res))
+
+if __name__ == '__main__':
+ unittest.main()

torch/_dynamo/utils.py

@@ -1409,8 +1409,12 @@ def to_tensor(t):
  ref = ref.to(res.dtype)
 
  # First try usual allclose
- if torch.allclose(ref, res, atol=tol, rtol=tol, equal_nan=equal_nan):
- return True
+ try:
+ if torch.allclose(ref, res, atol=tol, rtol=tol, equal_nan=equal_nan):
+ return True
+ except Exception:
+ logging.exception("An unexpected error occurred while comparing tensors with torch.allclose")
+ raise
 
  # Check error from fp64 version
  if fp64_ref.dtype == torch.float64: