Hi, you could use this functionality - torch.cdist — PyTorch 2.2 documentation which calculates the distance between each vector in ‘b’ to each vector in ‘a’.
E.g. torch.cdist(b, a, p=2) will give (128, 1600) shape, where each row represents the L2-norm distances between the corresponding value in ‘b’ and all values in ‘a’.