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