To install specific CUDA version builds for PyTorch, use
$ pip install torch -f https://download.pytorch.org/whl/cu101/torch_stable.html
cu101
corresponds to CUDA 10.1, and needs to be changed appropriately.
For , the pairwise distance matrix between each pair of these batched matrices is , where ... represent arbitrary batch dimension (think batches of pairs of and samples of dimension ).
def pairwise_dist(x, y):xx = (x * x).sum(dim=-1).unsqueeze(-1)yy = (y * y).sum(dim=-1).unsqueeze(-2)xy = torch.einsum('...ji,...ki->...jk', x, y)d = xx + yy - 2. * xyreturn d