def run_optimizer(self, loss):
    'Runs backwards pass and update the optimizer'
    self.check_inf_nan(loss)
    update_idx = (self.num_updates // self.get_global_batchsize())
    do_zero_grad = ((update_idx % self.optimizer_period) == 0)
    do_step = self._should_do_step()
    if do_zero_grad:
        self.optimizer.zero_grad()
    if (self.amp_type == AmpType.APEX):
        with apex.amp.scale_loss(loss, self.optimizer.optimizer) as scaled_loss:
            scaled_loss.backward()
    elif (self.amp_type == AmpType.PYTORCH):
        self.amp_grad_scaler.scale(loss).backward()
    else:
        loss.backward()
    if do_step:
        if (self.optimizer_period != 1):
            self._rescale_gradients((1 / self.optimizer_period))
        if (self.clip_grad_norm is not None):
            self._clip_gradients(self.clip_grad_norm)
        if (self.amp_type == AmpType.PYTORCH):
            self.amp_grad_scaler.step(self.optimizer, where=self.where)
            self.amp_grad_scaler.update()
        else:
            self.optimizer.step(where=self.where)
Runs backwards pass and update the optimizer
def run_optimizer(self, loss):
    
    self.check_inf_nan(loss)
    update_idx = (self.num_updates // self.get_global_batchsize())
    do_zero_grad = ((update_idx % self.optimizer_period) == 0)
    do_step = self._should_do_step()
    if do_zero_grad:
        self.optimizer.zero_grad()
    if (self.amp_type == AmpType.APEX):
        with apex.amp.scale_loss(loss, self.optimizer.optimizer) as scaled_loss:
            scaled_loss.backward()
    elif (self.amp_type == AmpType.PYTORCH):
        self.amp_grad_scaler.scale(loss).backward()
    else:
        loss.backward()
    if do_step:
        if (self.optimizer_period != 1):
            self._rescale_gradients((1 / self.optimizer_period))
        if (self.clip_grad_norm is not None):
            self._clip_gradients(self.clip_grad_norm)
        if (self.amp_type == AmpType.PYTORCH):
            self.amp_grad_scaler.step(self.optimizer, where=self.where)
            self.amp_grad_scaler.update()
        else:
            self.optimizer.step(where=self.where)