Scaling variables in manual_backward when using mixed precision

[marcmk6]

Hi,

The algorithm that I'm working on directly issues gradient of "loss" w.r.t. model predictions, so I need get related variables (model_prediction and gradient in the following example) scaled when using mixed precision. How can I achieve this?

def training_step(...):
    opt = self.optimizers()
    loss = torch.tensor(0) # not used in my case
    opt.zero_grad()
    # automatically applies scaling, etc...
    self.manual_backward(loss, {'model_prediction': model_prediction, 'grad': gradient})
    opt.step()

def backward(loss, optimizer, optimizer_idx, *args, **kwargs):
    # `loss` is not used
    prediction = kwargs['model_prediction']
    grad = kwargs['grad']
    # how to get these two variables scaled? 
    prediction.backward(grad)

Thanks

[marcmk6]

@rohitgr7 Could you help with this?