enable_python_dispatcher() in propagate_types (#191)

Author: laithsakka

helion/_compiler/type_propagation.py

@@ -17,6 +17,7 @@
 
 import sympy
 import torch
+from torch._dispatch.python import enable_python_dispatcher
 from torch._dynamo.convert_frame import compile_lock
 from torch.fx.experimental import proxy_tensor
 from torch.utils._pytree import tree_map_only
@@ -2151,7 +2152,7 @@ def _to_proxy(arg: TypeInfo) -> object:
 
 def propagate_types(func: HostFunction, fake_args: list[object]) -> None:
     # Lock needed since patch.object(torch.SymInt.__index__, ...) is not thread safe
-    with compile_lock, func:
+    with compile_lock, func, enable_python_dispatcher():
         global_scope = GlobalScope(function=func)
         local_scope = LocalScope(parent=global_scope)
         params = inspect.signature(func.fn).bind(*fake_args)