pytorch/c10/core/InferenceMode.cpp
Ailing Zhang e554731b32 Hide set_enabled since it's not public facing. (#58078)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58078

Test Plan: Imported from OSS

Reviewed By: albanD

Differential Revision: D28362048

Pulled By: ailzhang

fbshipit-source-id: 4c78a7c58860ec4963bc8d05d133ea26e47dcf00
2021-05-12 22:52:17 -07:00

21 lines
627 B
C++

#include <c10/core/InferenceMode.h>
#include <stdexcept>
namespace c10 {
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
thread_local bool InferenceMode_enabled = false;
// Invariant:
// is_enabled() ==
// !c10::impl::tls_is_dispatch_key_included(DispatchKey::ADInplaceOrView);
// InferenceMode::is_enabled() is in perf critical path (TensorImpl constructor)
// so it worths a separate TLS to skip the DispatchKeySet check.
bool InferenceMode::is_enabled() {
return InferenceMode_enabled;
}
void InferenceMode::_set_enabled(bool enabled) {
InferenceMode_enabled = enabled;
}
} // namespace c10