Should automatically pop modes

Summary:
Fixes pytorch/pytorch#108282

X-link: pytorch/pytorch#109157

Reviewed By: zou3519

Differential Revision: D49359181

Pulled By: yanboliang

fbshipit-source-id: f91d4aec67db980f3c66b7e0232106b89e34aef3

Author: yanboliang

exir/delegate.py

@@ -114,18 +114,16 @@ def fake_requires_grad(var):
 
 @executorch_call_delegate.py_impl(ProxyTorchDispatchMode)
 # pyre-ignore
-def call_delegate_proxy_torch_dispatch_mode(lowered_module, *args):
-    mode = _get_current_dispatch_mode()
-    assert mode is not None, "Mode should always be enabled for python fallback key"
-    with _pop_mode_temporarily() as mode:
-        res = trace_call_delegate(mode, executorch_call_delegate, lowered_module, *args)
+def call_delegate_proxy_torch_dispatch_mode(mode, lowered_module, *args):
+    res = trace_call_delegate(mode, executorch_call_delegate, lowered_module, *args)
     return res
 
 
 @executorch_call_delegate.py_impl(FakeTensorMode)
 # pyre-ignore
-def call_delegate_fake_tensor_mode(lowered_module, *args):
-    return lowered_module.original_module(*args)
+def call_delegate_fake_tensor_mode(mode, lowered_module, *args):
+    with mode:
+        return lowered_module.original_module(*args)
 
 
 @executorch_call_delegate.py_impl(torch._C.DispatchKey.Functionalize)