[https://nvbugs/5625972][fix] Add context manager to fix FakeTensorProp (#9047)

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>

Author: Fridah-nv

tensorrt_llm/_torch/auto_deploy/utils/_graph.py

@@ -7,6 +7,7 @@
 import torch.nn as nn
 import torch.utils._pytree as pytree
 from torch import fx
+from torch._dispatch.python import enable_python_dispatcher
 from torch._export.utils import _detect_fake_mode_from_gm
 from torch._prims_common import DeviceLikeType
 from torch._subclasses import FakeTensor, FakeTensorMode
@@ -210,7 +211,8 @@ def _run_shape_prop_single_gm(
 
     # run shape propagation if we have all the fake tensors
     if all(inp is not _NO_VAL for inp in inps):
-        FakeTensorProp(gm, fake_mode).propagate(*inps)
+        with enable_python_dispatcher():
+            FakeTensorProp(gm, fake_mode).propagate(*inps)
     else:
         ad_logger.warning("No fake tensors and no args available for shape propagation")
 

tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_build_small_single.py

@@ -114,15 +114,14 @@ def _check_ad_config(experiment_config: ExperimentConfig, llm_args: LlmArgs):
                 },
             },
         ),
-        pytest.param(
+        (
             "meta-llama/Llama-4-Scout-17B-16E-Instruct",
             {
                 "transforms": {
                     "insert_cached_attention": {"backend": "flashinfer"},
                     "compile_model": {"backend": "torch-opt"},
                 },
             },
-            marks=pytest.mark.skip(reason="https://nvbugs/5625972"),
         ),
         (
             "meta-llama/Llama-4-Scout-17B-16E-Instruct",