def METHOD_NAME(self, batch, batch_idx):
    self.seen_train_batches.append(batch)
    should_signal = self.trainer.fit_loop.epoch_loop._is_training_done if self.on_last_batch else batch_idx == 2
    if self.failure_on_step and self.failure_on_training and should_signal:
        self._signal()
    return super().METHOD_NAME(batch, batch_idx)