Gradient Penalty and AMP

[FeryET]

I am trying to implement a gradient penalty for a GAN implementation. From looking at Pytorch's documentation, it seems that the gradients should be computed on scaled inputs.

Should I implement a completely manual loop for this to be implemented, or can I implement it in the training_step with manual_backward. How can I access the gradient scaler to scale the input tensors.

Thanks!