commit 3d55d84ec2271b58e731f180788c3e564fc69cc8
author Jiashen Cao <jiashenc@meta.com> Wed Jun 12 01:52:09 2024 +0000
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Wed Jun 12 01:52:09 2024 +0000
tree abe8c877d60a02f9f1ea9c8bbaee9a4a5457da19
parent bb2a9955297fe064b99308a64a1ac43ab1a212c8

[Fix] Check tensor dtype before using torch.allclose in _trace log (#128438)

#### Issue
`torch.allclose` errors out during logging due to different dtypes.

#### Test
* `pytest test/test_jit.py`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/128438
Approved by: https://github.com/angelayi

torch/jit/_trace.py

diff --git a/torch/jit/_trace.py b/torch/jit/_trace.py
index 5bdd71f..7db8560 100644
--- a/torch/jit/_trace.py
+++ b/torch/jit/_trace.py
@@ -656,6 +656,8 @@
             return False
 
         if isinstance(orig, torch.Tensor):
+            if orig.dtype != loaded.dtype:
+                return False
             if not torch.allclose(orig, loaded):
                 return False
         else: