Hi there,
Have a question regarding how to leverage torch for general tensor operations (e.g., matmul, cdist, etc.) other than deep learning. So I do not have a training process but a simple calculation.
For instance, I would like to calculate the pairwise distance of two large matrices (100,000 samples, 128 dimensions) with four GPUs (cuda:0,1,2,3). A single GPU does not have enough memory for doing so.

A = torch.randn(100000, 128).cuda()
B = torch.randn(100000, 128).cuda()
pdist = torch.nn.PairwiseDistance(p=2)
pairwise_distance = pdist(A, B)

My questions are:

how to easily split the task to multiple GPUs (just like joblib with native Python)?
how to do the calculation in parallel (I could retrieve the gpu id and split the matrix into multiple fold with a for loop)
does pytorch multiprocessing also handle data split with multiple GPU? I am afraid that is not the case.

Thanks for the help and happy new year!