/
vshmidt
/
pytorch
Обзор
Документация
Войти
/
vshmidt
/
pytorch
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
caffe2/utils/threadpool/thread_pool_guard.h
23 строки
567 B
Akshit Khurana
Add guard to run on current thread (#52361)
03 мар 2021, 22:43
03 мар 2021, 22:43
ecd8e4c
Код
Авторство
О чём код?
#pragma once #include <c10/macros/Macros.h> namespace caffe2 { // A RAII, thread local (!) guard that enables or disables grad mode upon // construction, and sets it back to the original value upon destruction. struct TORCH_API _NoPThreadPoolGuard { static bool is_enabled(); static void set_enabled(bool enabled); _NoPThreadPoolGuard(): prev_mode_(_NoPThreadPoolGuard::is_enabled()) { _NoPThreadPoolGuard::set_enabled(true); } ~_NoPThreadPoolGuard() { _NoPThreadPoolGuard::set_enabled(prev_mode_); } private: bool prev_mode_; }; }