# Apply multi-head attention
output = attention.forward(X, Y, Y, valid_lens)