Gradient Overflow Until ZeroFloatDivision on WGAN-GP

Hello everyone,
I am trying to train a WGAN-GP model on O1 and O2 opt levels, but on the gradient penalty phase, I get gradient overflow error till the division with zero exception with both O1 and O2. I searched for some WGAN-GP apex codes, which I find in here;

https://github.com/hukkelas/progan-pytorch/blob/master/src/models/loss.py
using another scaled loss and scaling factor for GP. However, this also gives the same error.

My code for GP calculation as follows,

def _gradient_penalty(self, real_data, generated_data, gp_weight):

    batch_size = real_data.size()[0]

    # Calculate interpolation
    alpha = torch.rand(batch_size, 1, 1)
    alpha = alpha.expand_as(real_data)
    alpha.to(real_data.dtype)

    interpolated = alpha * real_data.data + (1 - alpha) * generated_data.data
    interpolated = Variable(interpolated, requires_grad=True)

    interpolated.to(real_data.dtype)
    # Calculate probability of interpolated examples
    prob_interpolated = self.discriminator(interpolated)

    # Calculate gradients of probabilities with respect to examples
    gradients = torch_grad(outputs=prob_interpolated, inputs=interpolated,
                           grad_outputs=torch.ones(prob_interpolated.size()),
                           create_graph=True, retain_graph=True)[0]

    gradients = gradients.view(gradients.size(0), -1)
    gradient_norm = gradients.norm(2, dim=1)
    gradient_penalty = ((gradient_norm - 1) ** 2).mean()
    return gp_weight * gradient_penalty`

       gradient_penalty = self._gradient_penalty(x, generated_data, gp_weight)
       d_gen = self.discriminator(generated_data)
       d_real = self.discriminator(x)
       d_loss = d_gen.mean() - d_real.mean() + gradient_penalty
       with amp.scale_loss(d_loss, self.disc_opt, loss_id=1) as scaled_loss:
           scaled_loss.backward()

Can you help me about that please?
cc @mcarilli

Reproduce the WGAN-GP gradient-penalty path shown in the issue with Apex O1 and O2, starting at _gradient_penalty and amp.scale_loss(...). Compare the overflow and zero-division behavior around torch_grad and the scaled discriminator loss. Done means identifying a supported correction or a confirmed Apex limitation, with a minimal reproduction and regression coverage if the repository has a suitable test location.