I have a batch = 256 of pair of vectors of dimension 32. The shape is so [256, 2, 32].
What I would want would be to perform the dot product between each of the pair of vectors of the batch, so that I end up with a tensor of shape [256, 32].
I don’t understand how you wish to end up with such dimensions. If one vector is X[:,0,:] and another is X[:,1,:], and you want to dot product them, the result should be a either a scalar, or a vector of length 256 or a vector of length 32 (if you want to perform dot product in one dimension).
Maybe you want elementwise product?
In that case it can be achieved by torch.prod(X, dim=1)