Multi-head attention as a PyTorch module
Design a basic multi-head attention block in PyTorch. What tensors do Q, K and V have, when do you split heads, and what common shape mistakes should you avoid?
Answer as you would in an interview
Answer aloud or write down a few key points. The text field is optional; the review opens after your attempt.