File size: 743 Bytes
ad5f26a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
# Config options to enable/disable C++ kernel for nn.functional.MHA
# and nn.TransformerEncoder
import torch
_is_fastpath_enabled: bool = True
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
def set_fastpath_enabled(value: bool) -> None:
"""Sets whether fast path is enabled"""
global _is_fastpath_enabled
_is_fastpath_enabled = value
|