torch.backends.mha 的源代码

# 配置选项以启用/禁用 nn.functional.MHA 和 nn.TransformerEncoder 的 C++ 内核 # 导入 torch _is_fastpath_enabled: bool = True
[文档]def get_fastpath_enabled() -> bool:
返回 TransformerEncoder 和 MultiHeadAttention 的快速路径是否可用
已启用，或如果 jit 是脚本化则显示为``True``。

.. 注意：
快速路径可能不会运行，即使 `get_fastpath_enabled` 返回
除非所有输入条件都满足，否则为 ``True``。
"""
如果没有调用 torch.jit.is_scripting()
返回 _is_fastpath_enabled
return True


[docs]def set_fastpath_enabled(value: bool) -> None:
设置是否启用快速路径
global _is_fastpath_enabled
_is_fastpath_enabled = value