def METHOD_NAME(
    self,
    trainer: "pl.Trainer",
    pl_module: "pl.LightningModule",
    batch: Any,
    batch_idx: int,
    dataloader_idx: int = 0,
) -> None:
    self._get_and_log_device_stats(trainer, "on_test_batch_start")