Thank you very much!
I’m little confused about the line:
conv = nn.Conv1d(1, num_heads, head_dim, head_dim) 

As far as I read the pytorch documentation, the nn.Conv1d should take the 4th argument as stride = head_dim. Do you explicitly left it as here or is it just a typo?
By the way, if I have a audio segment with N frames with 12 dimensions of each, should I modify your code this way:
input_dim = 12
frame_len = N
input = torch.rand(1, frame_len, input_dim)
...