import cupy
import ot.gpu

with Pool(processes = workers) as p_kappa: 
            W = list(p_kappa.imap(partial(W_comp, mx_all, dist, lamb), G.edges()),\
                          total = len(G.edges))

def W_comp(mx_all, dist, lamb,  e):
    i = e[0]
    j = e[1]

    W = np.zeros(5)
    for it in range(5):

        Nx = np.array(mx_all[i][1][it]).flatten()
        mx = mx_all[i][0][it].toarray().flatten()
        my = mx_all[j][0][it].toarray().flatten()

        dNx = dist[Nx,:].copy(order='C')

        W[it] = ot.sinkhorn(mx, my, dNx, lamb)