1 | #pragma once |
---|---|
2 | #include <ATen/core/TorchDispatchUtils.h> |
3 | |
4 | namespace at { |
5 | namespace impl { |
6 | |
7 | struct TORCH_API RestorePythonTLSSnapshot { |
8 | RestorePythonTLSSnapshot(); |
9 | ~RestorePythonTLSSnapshot(); |
10 | |
11 | private: |
12 | c10::impl::LocalDispatchKeySet saved_; |
13 | c10::impl::ForceDispatchKeyGuard guard_; |
14 | }; |
15 | |
16 | |
17 | // RAII guard to make working with the above TLS safer. |
18 | struct TORCH_API MaybeSetTLSOnEntryGuard { |
19 | public: |
20 | MaybeSetTLSOnEntryGuard(); |
21 | ~MaybeSetTLSOnEntryGuard(); |
22 | |
23 | private: |
24 | bool value_set_; |
25 | }; |
26 | |
27 | } // namespace impl |
28 | } // namespace at |
29 |