It’s quite complex (here 5) so I can’t really paste it all but for some reason the culprit seems to be changing requires_grad back and forth for the discriminator
# ------------------------------
# -----  Update Generator  -----
# ------------------------------
if self.d_opt is not None:
    for param in self.D.parameters():
        # continue
        param.requires_grad = False

self.update_G(batch)

# ----------------------------------
# -----  Update Discriminator  -----
# ----------------------------------

# unfreeze params of the discriminator
for param in self.D.parameters():
    # continue
    param.requires_grad = True

self.update_D(batch)

The error disappears if I comment-in continue or  equivalently if I comment our the 2 for loops around the gradient
Where both self.update_X(batch) for being either the generator (g) or the discriminator (d) methods are structured as:
with autocast():
    x_loss = self.get_x_loss(batch)
self.grad_scaler_x.scale(x_loss).backward()
self.grad_scaler_x.step(self.x_opt)
self.grad_scaler_x.update()

In both cases x_opt is a regular torch.optim.Adam(X.parameters())