mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
66a2600b6a
Summary: This is a re-attempt to land the iwyu header changes, by taking the diff from [PR 100304](https://github.com/pytorch/pytorch/pull/100304), and adding the bare minimal changes to make the diff build corectly in the internal builds. X-link: https://github.com/facebookresearch/pytorch3d/pull/1541 X-link: https://github.com/fairinternal/pytorch3d/pull/44 - Re-work D45769819 to fix header inclusions in c10 Test Plan: ``` buck2 build --no-remote-cache mode/dev-nosan //caffe2/c10/... buck2 build --no-remote-cache mode/dev-nosan //deeplearning/fbgemm/fbgemm_gpu/... buck2 build mode/dev-nosan //vision/fair/pytorch3d/pytorch3d:_C ``` Reviewed By: malfet Differential Revision: D45920611 Pull Request resolved: https://github.com/pytorch/pytorch/pull/101846 Approved by: https://github.com/malfet, https://github.com/Skylion007
45 lines
1.2 KiB
C++
45 lines
1.2 KiB
C++
#pragma once
|
|
|
|
#include <c10/core/AutogradState.h>
|
|
#include <c10/macros/Export.h>
|
|
|
|
namespace c10 {
|
|
|
|
struct C10_API GradMode {
|
|
static bool is_enabled();
|
|
static void set_enabled(bool enabled);
|
|
};
|
|
|
|
// A RAII, thread local (!) guard that enables or disables grad mode upon
|
|
// construction, and sets it back to the original value upon destruction.
|
|
struct C10_API AutoGradMode {
|
|
AutoGradMode(bool enabled) : prev_mode(GradMode::is_enabled()) {
|
|
GradMode::set_enabled(enabled);
|
|
}
|
|
~AutoGradMode() {
|
|
GradMode::set_enabled(prev_mode);
|
|
}
|
|
bool prev_mode;
|
|
};
|
|
|
|
// A RAII, thread local (!) guard that stops future operations from building
|
|
// gradients.
|
|
struct C10_API NoGradGuard : public AutoGradMode {
|
|
NoGradGuard() : AutoGradMode(/*enabled=*/false) {}
|
|
};
|
|
|
|
// A RAII, thread local (!) guard that enables or disables forward grad mode
|
|
// upon construction, and sets it back to the original value upon destruction.
|
|
struct C10_API AutoFwGradMode {
|
|
AutoFwGradMode(bool enabled)
|
|
: prev_mode(AutogradState::get_tls_state().get_fw_grad_mode()) {
|
|
AutogradState::get_tls_state().set_fw_grad_mode(enabled);
|
|
}
|
|
~AutoFwGradMode() {
|
|
AutogradState::get_tls_state().set_fw_grad_mode(prev_mode);
|
|
}
|
|
bool prev_mode;
|
|
};
|
|
|
|
} // namespace c10
|