mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
27c7158166
Summary: There is a module called `2to3` which you can target for future specifically to remove these, the directory of `caffe2` has the most redundant imports: ```2to3 -f future -w caffe2``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/45033 Reviewed By: seemethere Differential Revision: D23808648 Pulled By: bugra fbshipit-source-id: 38971900f0fe43ab44a9168e57f2307580d36a38
36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
|
|
|
|
|
|
import contextlib
|
|
import copy
|
|
import threading
|
|
|
|
_threadlocal_scope = threading.local()
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def arg_scope(single_helper_or_list, **kwargs):
|
|
global _threadlocal_scope
|
|
if not isinstance(single_helper_or_list, list):
|
|
assert callable(single_helper_or_list), \
|
|
"arg_scope is only supporting single or a list of helper functions."
|
|
single_helper_or_list = [single_helper_or_list]
|
|
old_scope = copy.deepcopy(get_current_scope())
|
|
for helper in single_helper_or_list:
|
|
assert callable(helper), \
|
|
"arg_scope is only supporting a list of callable helper functions."
|
|
helper_key = helper.__name__
|
|
if helper_key not in old_scope:
|
|
_threadlocal_scope.current_scope[helper_key] = {}
|
|
_threadlocal_scope.current_scope[helper_key].update(kwargs)
|
|
|
|
yield
|
|
_threadlocal_scope.current_scope = old_scope
|
|
|
|
|
|
def get_current_scope():
|
|
global _threadlocal_scope
|
|
if not hasattr(_threadlocal_scope, "current_scope"):
|
|
_threadlocal_scope.current_scope = {}
|
|
return _threadlocal_scope.current_scope
|