Here is a more minimal example. No more cuda/data parallel/nn.Module and such 
Still looking into why it fails
import torch
from torch import nn, cuda
from torch.autograd import Variable, grad
from torch.nn import functional as F

# Debug stuff
import torchviz
torch.autograd.set_detect_anomaly(True)

inputs = torch.ones((1, 3, 256, 256), requires_grad=True)

tmp1 = F.instance_norm(inputs)
tmp2 = F.threshold(tmp1, 0., 0., True)
prob_interpolated = torch.sigmoid(tmp2)

gradients = grad(outputs=prob_interpolated, inputs=inputs,
                 grad_outputs=torch.ones(prob_interpolated.size()),
                 create_graph=True, retain_graph=True)[0]

gradient_penalty = gradients.sum()

# Debug graph
torchviz.make_dot(gradient_penalty).view()
gradient_penalty.backward()