pytorch
diff --git a/‎torch/fx/experimental/proxy_tensor.py
Copy file name to clipboardExpand all lines: torch/fx/experimental/proxy_tensor.py
+3Lines changed: 3 additions & 0 deletions b/‎torch/fx/experimental/proxy_tensor.py
Copy file name to clipboardExpand all lines: torch/fx/experimental/proxy_tensor.py
+3Lines changed: 3 additions & 0 deletions
@@ -1136,13 +1136,16 @@ def _should_save_eager_input_vals(
     target: Any,
     args_kwargs: Optional[tuple[tuple[Argument, ...], dict[str, Argument]]] = None,
 ) -> bool:
+    from torch._higher_order_ops.invoke_subgraph import InvokeSubgraphHOP
+
     if not callable(target):
         return False
     if isinstance(
         target,
         (
             torch._higher_order_ops.triton_kernel_wrap.TritonKernelWrapperFunctional,
             torch._higher_order_ops.triton_kernel_wrap.TritonKernelWrapperMutation,
+            InvokeSubgraphHOP,
         ),
     ):
         return True