def METHOD_NAME(self):
    class _TrainerTestWrapper(training.MirroredStrategyTrainer):
        def __init__(self, test_instance, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.test_instance = test_instance
            self.losses = []
            self.sample_sizes = []
        def _accumulate_loss(self, loss, sample_size):
            self.losses.append(loss)
            self.sample_sizes.append(sample_size)
            return super()._accumulate_loss(loss, sample_size)
        def _apply_gradients(self, sample_size):
            self.test_instance.assertLen(self.losses, 4)
            self.test_instance.assertLen(self.sample_sizes, 4)
            self.test_instance.assertEqual(sample_size, sum(self.sample_sizes))
            return super()._apply_gradients(sample_size)
    model = _make_seq2seq_model(self.get_temp_dir())
    optimizer = model.get_optimizer()
    devices = tf.config.list_logical_devices(device_type="CPU")
    source_path = self._makeTextFile(
        "source.txt", ["1 2 3", "1 2 3 4 5", "1 2 3 4", "1 2 3 4 5 6 7"]
    )
    target_path = self._makeTextFile(
        "target.txt", ["1 2 3 4", "1 2 3", "1 2 3 4 5 6", "1 2 3 4 5"]
    )
    dataset = model.examples_inputter.make_training_dataset(
        source_path, target_path, batch_size=1, single_pass=True
    )
    trainer = _TrainerTestWrapper(self, model, optimizer, devices=devices)
    tf.config.run_functions_eagerly(True)
    summary = trainer(dataset, accum_steps=2)
    tf.config.run_functions_eagerly(False)
    self.assertEqual(summary["num_steps"], 1)
    self.assertAllClose(
        summary["last_loss"], sum(trainer.losses) / sum(trainer.sample_sizes)
    )