Code: self._optimizer_step_flag is True every 4th step. GA = 4
def _fp16_backwards(self,loss):

    self.logger.debug(f'loss scale = {_amp_state.loss_scalers[0].loss_scale()}, loss = {loss.item()}')

    with amp.scale_loss(loss,self.optimizer, delay_unscale= not self._optimizer_step_flag) as scaled_loss:

        self.logger.debug(f'scaled loss = {scaled_loss.item()}')

        scaled_loss.backward()

    # Clip gradients only when optimizer is about to step. Fixes divergence issue. 

    # TODO: patch _optimize and move code there.

    # TODO: clip grads for non FP16 training also

    if self._optimizer_step_flag:

        print(f'optimizer norm = {get_norm(amp.master_params(self.optimizer))}')

        print(f'model , optimizer max grad before clipping \

            {get_max_grad(self.model.parameters())}, {get_max_grad(amp.master_params(self.optimizer))}')

        torch.nn.utils.clip_grad_norm_(amp.master_params(self.optimizer), 1.0)

        print(f'model , optimizer max grad after clipping \

            {get_max_grad(self.model.parameters())}, {get_max_grad(amp.master_params(self.optimizer))}')

        print(f'max optimizer parameter : {max([torch.max(p).item() for p in amp.master_params(self.optimizer)])}')