From 3c1dd4e752162db1d7cc993187da623cec3f0d2b Mon Sep 17 00:00:00 2001
From: Horace He <chilli@fb.com>
Date: Thu, 31 Mar 2022 01:32:34 +0000
Subject: [PATCH] Removed python dispatch keys from dispatch key extraction

Fixes this issue

https://github.com/facebookresearch/torchdynamo/issues/81
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74971
Approved by: https://github.com/zou3519
---
 c10/core/DispatchKeySet.h | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/c10/core/DispatchKeySet.h b/c10/core/DispatchKeySet.h
index 79d39652219..63c48f04056 100644
--- a/c10/core/DispatchKeySet.h
+++ b/c10/core/DispatchKeySet.h
@@ -329,7 +329,8 @@ static inline DispatchKey legacyExtractDispatchKey(DispatchKeySet s) {
   // here.  At the moment, autograd keys and ADInplaceOrView key need this
   // treatment;
   return (s - autograd_dispatch_keyset_with_ADInplaceOrView -
-          autocast_dispatch_keyset)
+          autocast_dispatch_keyset -
+          DispatchKeySet({DispatchKey::PythonTLSSnapshot, DispatchKey::Python}))
       .highestPriorityTypeId();
 }