PyTorch How to code Multi Head Self Attention in parallel?
I want to encode the word (embedding) sequence with 16-Head Self-Attention.
Currently I use a nn.ModuleList together with a for loop to generate the output of each head then concatenate all of them. This approach is extremely slow and I wonder if there's way to code MHA in parallel?
To generalize, I would like to know if I can 'stack' multiple nn.Linear, as I feed the input vector, the multiple outputs will be computed in parallel.

I figured it out. Since nn.Linear is acctually an affine transformation with a weights matrix and a bias matrix, one can easily wrap such matrices in nn.Parameter and take advantage of broadcast semantics to achieve the goal.

Edit: I also find a nn.Linear(d_model, n_heads*d_key) functions identically.
