I want to compute gradients in the standard manner:
grad = torch.autograd.grad(loss, model.parameters()

I want to use AMP. To get the scaled loss, apex api suggests:
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

But I don’t want to call loss.backward() directly. I compute the gradients in the manner I showed above
I tried this:
with amp.scale_loss(loss, optimizer) as scaled_loss:
    grad = torch.autograd.grad(scaled_lossloss, model.parameters()

This makes the loss nan. How to compute gradients with apex ?