Implement parallelization over heads #2
Description
Currently the attention over heads runs in serial:
Line 101 in 01eb84b
We should try to parallelize it and see if we get any speedups.
{{ message }}
Currently the attention over heads runs in serial:
Line 101 in 01eb84b
We should try to parallelize it and see if we get any speedups.