fix(dtensor): support torch dispatch mode exclude

support torch dispatch python mode to tls excluded which will cause dispatch hang
pytorch · Apr 4, 2024 · 3bdfde2 · 3bdfde2
1 parent 86190b6
commit 3bdfde2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/cc/torchdistx/deferred_init.cc b/src/cc/torchdistx/deferred_init.cc
@@ -944,7 +944,7 @@ void runDeferredInitHandler(const OperatorHandle& op, DispatchKeySet key_set, St
 }
 
 void enableDeferredInitHandler(bool value) noexcept {
- c10::impl::tls_set_dispatch_key_included(DispatchKey::DeferredInit, value);
+ c10::impl::tls_set_dispatch_key_excluded(DispatchKey::DeferredInit, value);
 }
 
 bool isDeferredInitEnabled() noexcept {