Convert the gradients back to FP32 and divide them by the scale.

def on_backward_end(self, **kwargs:Any)->None:
        "Convert the gradients back to FP32 and divide them by the scale."
        if self.dynamic and grad_overflow(self.model_params) and self.loss_scale > 1:
            self.loss_scale /= 2
            self.noskip = 0
            #The step will be skipped since we don't update the master grads so they are all None or zero
        else:
            model_g2master_g(self.model_params, self.master_params, self.flat_master)
            for group in self.master_params:
                for param in group:
                    if param.grad is not None: param.grad.div_(self.loss_scale)
            if self.clip is not None:
                for group in self.master_params: nn.utils.clip_grad_norm_(group, self.clip)
            if not self.dynamic: return
            self.noskip += 1
            if self.noskip >= self.max_noskip and self.loss_scale < self.max_scale:
                self.loss_scale *= 2
                self.noskip = 0