torch.backends.mha — PyTorch 2.7 documentation (original) (raw)

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