Как сопоставить функцию Joblib Parallel с функцией DistributedDataParallel PyTorch

Question

У меня есть следующий код, который использует Parallel Joblib, и я хочу реализовать его в PyTorch и запустить его с помощью графических процессоров. Я читаю документацию PyTorch DistributedDataParallel , но, похоже, не могу понять этого.

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

Любая помощь будет принята с благодарностью. Заранее большое спасибо!