I am not using any wrong way to compute pow function.
#define loss function
def root(var):
    res=(var**0.5)
    return res

def reshape(tens):
    return torch.reshape(tens,(-1,))

def W2_dis(mean1, logvar1, mean2, logvar2):
    mean1=torch.reshape(mean1,(-1,))
    mean2=torch.reshape(mean1,(-1,))
    logvar1=torch.reshape(logvar1,(-1,))
    logvar2=torch.reshape(logvar2,(-1,))
    agg_var=torch.sum(logvar1.exp().pow(0.25)-logvar2.exp().pow(0.25))**2
    agg_mean=(torch.norm(mean1-mean2, 2))**2
    distance =(agg_mean+agg_var).pow(0.5)
    return distance

def wl1(m1, v1,m2, v2, m3, v3):
    ik=(-1*W2_dis(m1,v1,m3,v3)).exp()
    ij=W2_dis(m1,v1,m2,v2)
    r_loss=(ij**2+ik)
    return r_loss

def l2(x1,recon_x1,z1):
    x1=torch.reshape(x1,(-1,)).type(torch.cuda.FloatTensor)
    x1=x1/torch.sum(x1)
    #print(x1)
    recon_x1=torch.reshape(recon_x1,(-1,))
    bracket=torch.norm((x1*(x1 - recon_x1)),2)**2
    E_px=torch.mean(x1)
    E_qz=torch.mean(z1)
    loss=E_px*bracket*E_qz
    return loss

def wl2(x1,recon_x1,z1,x2,recon_x2,z2,x3,recon_x3,z3):
    loss1=l2(x1,recon_x1,z1)
    loss2=l2(x2,recon_x2,z2)
    loss3=l2(x3,recon_x3,z3)
    loss_final=torch.min(torch.min(loss1,loss2),torch.min(loss2,loss3))
    return loss_final

def wasserstein_loss(x1,recon_x1,z1, mu1, var1,x2,recon_x2,z2, mu2, var2,x3,recon_x3,z3, mu3, var3):
    l1_loss=wl1(mu1, var1, mu2, var2, mu3, var3)
    l2_loss=wl2(x1,recon_x1,z1,x2,recon_x2,z2,x3,recon_x3,z3)
    w_final=l1_loss+(0.4)*l2_loss
    return w_final

def rocauc_score(y_true,y_pred):
    y_true=torch.reshape(y_true,(-1,))
    y_pred=torch.reshape(y_pred,(-1,))
    y_true=y_true.detach().cpu().numpy()
    y_pred=y_pred.detach().cpu().numpy()
    #print(y_pred,y_true)
    return roc_auc_score(y_true,y_pred)

def wasserstein_acc(data1,recon_batch1,data2,recon_batch2, data3,recon_batch3):
    score1=rocauc_score(data1,recon_batch1)
    score2=rocauc_score(data2,recon_batch2)
    score3=rocauc_score(data3,recon_batch3)
    auroc=(score1+score2+score3)/3
    return auroc