def METHOD_NAME(
    self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", outputs: STEP_OUTPUT, batch: Any, batch_idx: int
) -> None:
    self._get_and_log_device_stats(trainer, "on_train_batch_end")