def test_batch_level_batch_indices(tmp_path):
    """Test that batch_indices are returned when `return_predictions=False`."""
    DummyPredictionWriter.write_on_batch_end = Mock()

    class CustomBoringModel(BoringModel):
        def on_predict_epoch_end(self, *args, **kwargs):
            assert self.trainer.predict_loop.epoch_batch_indices == [[]]

    writer = DummyPredictionWriter("batch")
    model = CustomBoringModel()
    dataloader = DataLoader(RandomDataset(32, 64), batch_size=4)
    trainer = Trainer(default_root_dir=tmp_path, logger=False, limit_predict_batches=4, callbacks=writer)
    trainer.predict(model, dataloaders=dataloader, return_predictions=False)

    writer.write_on_batch_end.assert_has_calls([
        call(trainer, model, ANY, [0, 1, 2, 3], ANY, 0, 0),
        call(trainer, model, ANY, [4, 5, 6, 7], ANY, 1, 0),
        call(trainer, model, ANY, [8, 9, 10, 11], ANY, 2, 0),
        call(trainer, model, ANY, [12, 13, 14, 15], ANY, 3, 0),
    ])
Test that batch_indices are returned when `return_predictions=False`.
tests/tests_pytorch/callbacks/test_prediction_writer.py