mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
f3fce597e9
See https://github.com/pytorch/pytorch/pull/129751#issue-2380881501. Most changes are auto-generated by linter. You can review these PRs via: ```bash git diff --ignore-all-space --ignore-blank-lines HEAD~1 ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/129769 Approved by: https://github.com/ezyang
33 lines
837 B
Python
33 lines
837 B
Python
# mypy: allow-untyped-defs
|
|
from contextlib import contextmanager
|
|
|
|
import torch
|
|
from torch.backends import __allow_nonbracketed_mutation, ContextProp, PropModule
|
|
|
|
|
|
__all__ = ["is_available", "flags", "set_flags"]
|
|
|
|
|
|
def is_available():
|
|
r"""Return whether PyTorch is built with NNPACK support."""
|
|
return torch._nnpack_available()
|
|
|
|
|
|
def set_flags(_enabled):
|
|
r"""Set if nnpack is enabled globally"""
|
|
orig_flags = (torch._C._get_nnpack_enabled(),)
|
|
torch._C._set_nnpack_enabled(_enabled)
|
|
return orig_flags
|
|
|
|
|
|
@contextmanager
|
|
def flags(enabled=False):
|
|
r"""Context manager for setting if nnpack is enabled globally"""
|
|
with __allow_nonbracketed_mutation():
|
|
orig_flags = set_flags(enabled)
|
|
try:
|
|
yield
|
|
finally:
|
|
with __allow_nonbracketed_mutation():
|
|
set_flags(orig_flags[0])
|