How to obtain the values of query and key in the attention module"
I believe you are talking about torch.nn.MultiheadAttention.
Multihead attention expects you to pass the query, key and value tensors in its forward pass. So, you would have access to them already.
I would be able to help you more if you could explain your question in a bit more detail.