/torch-multi-head-attention

Multi-head attention in PyTorch

Primary LanguagePythonMIT LicenseMIT

PyTorch Multi-Head Attention

Travis Coverage

Install

pip install torch-multi-head-attention

Usage

from torch_multi_head_attention import MultiHeadAttention

MultiHeadAttention(in_features=768, head_num=12)