commit a4075fca9adbfe3af0d7e7d393c99159c9440383
author Basil Hosmer <bhosmer@fb.com> Fri May 14 16:11:53 2021 -0700
committer Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> Fri May 14 16:12:57 2021 -0700
tree 2b20b09e0eebb63daa6ac4779a34fef69fbf808c
parent 3dc70d8f786afc5c150fc6f5c39f857eed244f2b

extract dispatch keys from optional Tensors (unboxed) (#58296)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58296

Test Plan: Imported from OSS

Reviewed By: Chillee

Differential Revision: D28436822

Pulled By: bhosmer

fbshipit-source-id: 8031c9a3c121483dd0e5ed7b8b165952477108e4

aten/src/ATen/core/dispatch/DispatchKeyExtractor.h

diff --git a/aten/src/ATen/core/dispatch/DispatchKeyExtractor.h b/aten/src/ATen/core/dispatch/DispatchKeyExtractor.h
index 7011686..9d3cc67 100644
--- a/aten/src/ATen/core/dispatch/DispatchKeyExtractor.h
+++ b/aten/src/ATen/core/dispatch/DispatchKeyExtractor.h
@@ -55,6 +55,11 @@
     void operator()(const at::Tensor& x) {
       ts = ts | x.key_set();
     }
+    void operator()(c10::optional<at::Tensor> x) {
+      if (x.has_value()) {
+        ts = ts | x->key_set();
+      }
+    }
     void operator()(at::ArrayRef<at::Tensor> xs) {
       for (const auto& x : xs) {
         ts = ts | x.key_set();