Source code for torch.backends.mha

# Config options to enable/disable C++ kernel for nn.functional.MHA
# and nn.TransformerEncoder
import torch

_is_fastpath_enabled: bool = True

[docs]def get_fastpath_enabled() -> bool: """Returns whether fast path for TransformerEncoder and MultiHeadAttention is enabled, or ``True`` if jit is scripting. ..note: The fastpath might not be run even if ``get_fastpath_enabled`` returns ``True`` unless all conditions on inputs are met. """ if not torch.jit.is_scripting(): return _is_fastpath_enabled return True
[docs]def set_fastpath_enabled(value: bool) -> None: """Sets whether fast path is enabled""" global _is_fastpath_enabled _is_fastpath_enabled = value


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources