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