code: def test_prediction_writer_batch_indices(num_workers):
      DummyPredictionWriter.write_on_batch_end = Mock()
      DummyPredictionWriter.write_on_epoch_end = Mock()
  
      dataloader = DataLoader(
          RandomDataset(32, 64), batch_size=4, num_workers=num_workers
      )
      model = BoringModel()
      writer = DummyPredictionWriter("batch_and_epoch")
      trainer = Trainer(limit_predict_batches=4, callbacks=writer)
      trainer.predict(model, dataloaders=dataloader)
  
      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),
          ]
      )
  
      writer.write_on_epoch_end.assert_has_calls(
          [
              call(
                  trainer,
                  model,
                  ANY,
                  [
                      [
                          [0, 1, 2, 3],
                          [4, 5, 6, 7],
                          [8, 9, 10, 11],
                          [12, 13, 14, 15],
                      ]
                  ],
              ),
          ]
      )
  
file_path: pytorch-lightning/tests/tests_pytorch/callbacks/test_prediction_writer.py
imports: from unittest.mock import ANY, Mock, call
  import pytest
  from lightning.pytorch import Trainer
  from lightning.pytorch.callbacks import BasePredictionWriter
  from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from torch.utils.data import DataLoader
label_desc: Use image node_id 3 for calling a global function with example usage: test_prediction_writer_batch_indices(num_workers) without return types