code: def training_step(self, batch: Batch, batch_idx):
      out = self.model(*batch.inputs)
      loss, num_frame = self.loss_fn(*out, *batch.labels)
      self.log_metric(batch, out, loss, "train")
  
      # normalize the loss based on the sum of num_frame across all GPUs
      num_frames = self.all_gather(num_frame)
      self.log(
          "Gathered number of frames",
          num_frames.float().sum(),
          on_step=True,
          on_epoch=True,
      )
      loss *= (
          num_frames.size(0) / num_frames.sum()
      )  # world size / num_frames
  
      return loss
  
file_path: audio/examples/self_supervised_learning/lightning_modules.py
imports: from collections import namedtuple
  from typing import Callable, Optional
  import lightning.pytorch
  import torch
  import torch.nn
  from torch.optim.optimizer import Optimizer
label_desc: Use image node_id 3 for calling the SSLPretrainModule obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss