John_J_Watson:

I tried using PairwiseDistance but this gives me a back a tensor of size 256, instead of 64,256.


This should normally return a [64] tensor or [64, 1] tensor if keepdim=True
t1 = torch.randn(64, 256)
t2 = torch.randn(64, 256)

t_dist = F.pairwise_distance(t1, t2)

print(t_dist.shape)