I’m trying to take advantage of Pytorch’s autograd feature and perform matrix-matrix multiplication $A \times B$ where matrix A is represented as a list of Tensors each on a separate GPU.
What is the best way of distributing this task across multiple GPUs and then collecting the results from each GPU onto one? It doesn’t seem to fit in with the paradigm of torch.nn.DataParallel where one model is replicated on each GPU and the data is passed through the model and then collected.