torch.backends.nnpack 的源代码

# mypy: 允许未类型注解的 from contextlib import contextmanager import torch from torch.backends import __allow_nonbracketed_mutation, ContextProp, PropModule __all__ = ["is_available", "flags", "set_flags"]
[文档]def is_available():
返回 PyTorch 是否带有 NNPACK 支持。
返回 torch._nnpack_available()


[文档]def set_flags(_enabled):
设置是否全局启用 nnpack。
orig_flags = (torch._C._get_nnpack_enabled(),)
torch._C._set_nnpack_enabled(_enabled)
return orig_flags


[文档]@contextmanager
def flags(enabled=False):
r"""全局启用 nnpack 的上下文管理器"""
with __allow_nonbracketed_mutation():
orig_flags = set_flags(enabled)
try:
yield
finally:
with __allow_nonbracketed_mutation():
设置标志（orig_flags[0]）