I have the following code below which uses Joblib’s Parallel and I want to implement this in PyTorch and run it with GPUs. I am reading through PyTorch’s DistributedDataParallel documentation 2, but I can’t seem to figure this out.
import numpy as np
import torch
from joblib import Parallel, delayed
from torch.nn.parallel import DistributedDataParallel as DDP

X = np.array([[1, 3, 2, 3], [2, 3, 5, 6], [1, 2, 3, 4]])
X = torch.DoubleTensor(X).cuda()

def X_power_func(j):
    X_power = X**j
    return X_power

results = Parallel(n_jobs = 4)(delayed(X_power_func)(j) for j in range(8))   # how do I map this to 
                                                                             # PyTorch's
                                                                             # DistributedDataParallel

Any help would really be appreciated. Many thanks in advance!