Here is a self-contained example, but as you can see it actually works now! This confirms that there is a bug somewhere in my code. I’ll take a closer look and report back once I find it:
import torch as tr
from torch.autograd import grad

# Generate data, 1000 samples from a normal distribution
X = tr.normal(tr.ones(1000)).view(-1, 1)
Z = tr.normal(tr.zeros(1000)).view(-1, 1)

# Define the generator and discriminator
# Generator: Az + b
class G(tr.nn.Module):
    def __init__(self):
        super(G, self).__init__()
        self.Ab = tr.nn.Linear(1, 1)
    
    def forward(self, x):
        return self.Ab(x)

# Discriminator: x^T C x + d^T x
class D(tr.nn.Module):
    def __init__(self):
        super(D, self).__init__()
        self.C = tr.nn.Bilinear(1, 1, 1, bias=False)
        self.d = tr.nn.Linear(1, 1, bias=False)
    
    def forward(self, x):
        return self.C(x, x) + self.d(x)

g = G()
d = D()

# Test to make sure outputs are valid
print("Generator Test:", g(Z).mean())
print("Discriminator Test:", d(X).mean(), d(g(Z)).mean())

# Define loss:
generator_loss = -1 * d(g(Z)).mean()
discriminator_loss = d(g(Z)).mean() - d(X).mean()

# Test to make sure loss is differentiable
print("Generator Loss:", generator_loss)
print("Discriminator Loss:", discriminator_loss)

# Get gradient wrt parameters
generator_gradient = grad(generator_loss, g.parameters(), retain_graph=True, create_graph=True)
discriminator_gradient = grad(discriminator_loss, d.parameters(), retain_graph=True, create_graph=True)

# Both tensors should still have grad_fn
print("Generator Gradient:", generator_gradient)
print("Discriminator Gradient:", discriminator_gradient)

# Take 2-norm of gradient components
generator_norm = tr.norm(tr.cat([tr.flatten(i) for i in generator_gradient]))
discriminator_norm = tr.norm(tr.cat([tr.flatten(i) for i in discriminator_gradient]))

# Add to loss function, should STILL have grad_fn
final_gen_loss = generator_loss + generator_norm
final_disc_loss = discriminator_loss + discriminator_norm

print("Final Generator Loss:", final_gen_loss)
print("Final Discriminator Loss:", final_disc_loss)

Output:
Generator Test: tensor(-0.4425, grad_fn=<MeanBackward0>)
Discriminator Test: tensor(-0.9662, grad_fn=<MeanBackward0>) tensor(0.3021, grad_fn=<MeanBackward0>)
Generator Loss: tensor(-0.3021, grad_fn=<MulBackward0>)
Discriminator Loss: tensor(1.2683, grad_fn=<SubBackward0>)
Generator Gradient: (tensor([[-0.0995]], grad_fn=<TBackward>), tensor([0.6430], grad_fn=<ViewBackward>))
Discriminator Gradient: (tensor([[[-1.6031]]], grad_fn=<AddBackward0>), tensor([[-1.3847]], grad_fn=<AddBackward0>))
Final Generator Loss: tensor(0.3486, grad_fn=<AddBackward0>)
Final Discriminator Loss: tensor(3.3866, grad_fn=<AddBackward0>)