def run_step_grad_acc(self, gradient_accumulation_steps=1):
        """
        Implement the AMP training logic.
        The batch at each step will be divided by this integer and
        gradient will be accumulated over gradient_accumulation_steps steps.
        """
        assert self._trainer.model.training, "[AMPTrainer] model was changed to eval mode!"
        assert torch.cuda.is_available(), "[AMPTrainer] CUDA is required for AMP training!"
        from torch.cuda.amp import autocast

        start = time.perf_counter()
        data = next(self._trainer._data_loader_iter)
        data_time = time.perf_counter() - start

        with autocast():
            loss_dict = self._trainer.model(data)
            if isinstance(loss_dict, torch.Tensor):
                loss_dict = loss_dict / gradient_accumulation_steps
                losses = loss_dict
                loss_dict = {"total_loss": loss_dict}
            else:
                losses = sum(loss_dict.values())
                losses = losses / gradient_accumulation_steps

        self._trainer.grad_scaler.scale(losses).backward()

        if (self._trainer.iter + 1) % gradient_accumulation_steps == 0:

            self._trainer._write_metrics(loss_dict, data_time)

            self._trainer.grad_scaler.step(self._trainer.optimizer)
            self._trainer.grad_scaler.update()
            self._trainer.optimizer.zero_grad()

        Implement the AMP training logic.
        The batch at each step will be divided by this integer and
        gradient will be accumulated over gradient_accumulation_steps steps.
        
layoutlmv3/examples/object_detection/ditod/mytrainer.py