mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
fd65bd755d
Something I've noticed is that a lot of the distributed sites don't render on our docs at all, but if they ever do, the notes will render properly now 😛 Pull Request resolved: https://github.com/pytorch/pytorch/pull/142868 Approved by: https://github.com/albanD
26 lines
718 B
Python
26 lines
718 B
Python
# 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
|