Thanks for replying @ptrblck . I am so sorry for not putting this in the post.
It is called via out = attention(a, b, a). The output is in the shape of the key and the value and the query are the same. I will append this in the original post too
Edit - Initially I had assumed PyTorch might be storing intermediate tensors for grad calculation. But even in a no_grad region, I observe the same behavior, the same 9.77Gb comes up(for a 8Gb GPU). Can anyone reproduce this??