I get Scatter is not differentiable twice when trying to backward the gradient penalty suggested in the WGAN paper. I’m running the latest version of pytorch and the models are using DataParallel!
    # discriminator
    d_real_out = disc(mixed, clean)
    d_fake = disc(fake)
    loss_real, loss_fake = torch.mean(d_fake), -torch.mean(d_real_out)
    grad_penalty = gradient_penalty(disc, mixed, clean, fake)
    loss_real.backward()
    loss_fake.backward()
    grad_penalty.backward()

    def compute_gradient_penalty(discriminator, mixed, real, fake, LAMBDA=10,
                                 use_cuda=True):
    alpha = torch.rand(real.size(0), 1)
    alpha = alpha.expand(real.size())
    alpha = alpha.contiguous()
    alpha = alpha.cuda() if use_cuda else alpha
    interpolates = alpha * real.data + ((1 - alpha) * fake.data)

    if use_cuda:
        interpolates = interpolates.cuda()
    interpolates = Variable(interpolates, requires_grad=True)
    disc_interpolates = discriminator(mixed, interpolates)

    gradients = torch.autograd.grad(
        outputs=disc_interpolates, inputs=interpolates,
        grad_outputs=torch.ones(disc_interpolates.size()).cuda() if use_cuda else torch.ones(disc_interpolates.size()),
        create_graph=True, retain_graph=True, only_inputs=True)[0]
    penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * LAMBDA
    return penalty