From 2289a12f21c54da93bf5d696e3f9aea83dd9c10d Mon Sep 17 00:00:00 2001 From: Horace He Date: Sat, 21 Aug 2021 17:13:27 -0700 Subject: [PATCH] Made FuncTorchBatched decompose CompositeImplicitAutograd (#63616) Summary: See https://github.com/facebookresearch/functorch/issues/56 Pull Request resolved: https://github.com/pytorch/pytorch/pull/63616 Reviewed By: zou3519 Differential Revision: D30438316 Pulled By: Chillee fbshipit-source-id: e84446d9f68b87daa0cfff75b3b8a972f36ec85a --- c10/core/DispatchKeySet.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/c10/core/DispatchKeySet.cpp b/c10/core/DispatchKeySet.cpp index 404acc7..21433d4 100644 --- a/c10/core/DispatchKeySet.cpp +++ b/c10/core/DispatchKeySet.cpp @@ -32,8 +32,8 @@ bool isBackendDispatchKey(DispatchKey t) { // math_dispatch_keyset contains all keys in backend_dispatch_keyset and // autograd_dispatch_keyset Alias key DispatchKey::CompositeImplicitAutograd // maps to math_dispatch_keyset. -constexpr DispatchKeySet math_dispatch_keyset = - backend_dispatch_keyset | autograd_dispatch_keyset; +constexpr DispatchKeySet math_dispatch_keyset = backend_dispatch_keyset | + autograd_dispatch_keyset | DispatchKeySet({DispatchKey::FuncTorchBatched}); DispatchKeySet getRuntimeDispatchKeySet(DispatchKey t) { TORCH_INTERNAL_ASSERT(t != DispatchKey::Undefined); -- 2.7.4