[easy][AOT] Fix shortcut path for simple tuple/list spec (#118460)

XuehaiPan · pytorchmergebot · commit 2728c9137dee · 2024-01-27T19:04:12.000Z
`type(self.spec)` is always `TreeSpec` and the condition is always `False`. This PR changes it to `self.spec.type`, which is the type of tree that the spec represents. Pull Request resolved: #118460 Approved by: https://github.com/Skylion007
diff --git a/torch/_functorch/_aot_autograd/utils.py b/torch/_functorch/_aot_autograd/utils.py
@@ -126,7 +126,7 @@ def set(self, spec: pytree.TreeSpec) -> None:
         assert self.spec is None or self.spec == spec
         assert spec is not None
         self.spec: pytree.TreeSpec = spec
-        if type(self.spec) in {tuple, list} and all(
+        if self.spec.type in {tuple, list} and all(
             child.is_leaf() for child in spec.children_specs
         ):
             self.is_simple = True
