1 #pragma once 2 3 #include <c10/macros/Macros.h> 4 5 namespace caffe2 { 6 7 // A RAII, thread local (!) guard that enables or disables grad mode upon 8 // construction, and sets it back to the original value upon destruction. 9 struct TORCH_API _NoPThreadPoolGuard { 10 static bool is_enabled(); 11 static void set_enabled(bool enabled); 12 _NoPThreadPoolGuard_NoPThreadPoolGuard13 _NoPThreadPoolGuard(): prev_mode_(_NoPThreadPoolGuard::is_enabled()) { 14 _NoPThreadPoolGuard::set_enabled(true); 15 } ~_NoPThreadPoolGuard_NoPThreadPoolGuard16 ~_NoPThreadPoolGuard() { 17 _NoPThreadPoolGuard::set_enabled(prev_mode_); 18 } 19 private: 20 bool prev_mode_; 21 }; 22 23 } 24