class_docstr: callback to be used with LightningTrainer and torchmetric
  Warning: 'optimizer.enable_state_sharding=True' is not supported
code: def on_train_batch_end(
      self,
      trainer: Trainer,
      pl_module: LightningModule,
      outputs: List,
      batch: SampleList,
      batch_idx: int,
      dataloader_idx: int,
  ):
      # prepare the next batch
      self.lightning_trainer.data_module.train_loader.change_dataloader()
      self.lightning_trainer.torchmetrics.update(batch, outputs)
      # log
      if (
          self._get_num_updates_for_logging(trainer)
          % self.trainer_config.log_every_n_steps
          == 0
      ):
          num_updates = self._get_num_updates_for_logging(trainer)
          max_updates = trainer.max_steps
          extra = self._get_train_extra_log(trainer, pl_module)
          self._log_metrics_and_extra(
              extra, num_updates, max_updates, log_type="train"
          )
          self.lightning_trainer.torchmetrics.reset()
          self.train_timer.reset()
  
imports: import logging
  from typing import Any, Dict, List, Optional
  import torch
  from mmf.common.registry import registry
  from mmf.common.sample import SampleList
  from mmf.utils.timer import Timer
  from pytorch_lightning import LightningModule, Trainer
  from pytorch_lightning.callbacks.base import Callback
label_desc: Use image node_id 3 for calling the LightningTorchMetricsCallback obj's underlying member method code with example usage: obj.on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) without return types