docstr: Custom training step with loss normalization and automatic mixed precision training.
  
  By default, DDP does the following on each train step:
  - For each GPU, compute loss and gradient on shard of training data.
  - Sync and average gradients across all GPUs. The final gradient
    is (sum of gradients across all GPUs) / N, where N is the world
    size (total number of GPUs).
  - Update parameters on each GPU.
  
  Here, we do the following:
  - For k-th GPU, compute loss and scale it by (N / num_frames), where num_frames is
    the sum of masked frames across all GPUs. Compute gradient from scaled loss.
  - Sync and average gradients across all GPUs. The final gradient
    is (sum of gradients across all GPUs) / num_frames.
  - Update parameters on each GPU.
  
  Doing so allows us to account for the variability in number of masked frames in
  variable-length sequential data.
docstr_tok: ["Custom","training","step","with","loss","normalization","and","automatic","mixed","precision","training",".","By","default",",","DDP","does","the","following","on","each","train","step",":","-","For","each","GPU",",","compute","loss","and","gradient","on","shard","of","training","data",".","-","Sync","and","average","gradients","across","all","GPUs",".","The","final","gradient","is","(","sum","of","gradients","across","all","GPUs",")","\/","N",",","where","N","is","the","world","size","(","total","number","of","GPUs",")",".","-","Update","parameters","on","each","GPU",".","Here",",","we","do","the","following",":","-","For","k-th","GPU",",","compute","loss","and","scale","it","by","(","N","\/","num_frames",")",",","where","num_frames","is","the","sum","of","masked","frames","across","all","GPUs",".","Compute","gradient","from","scaled","loss",".","-","Sync","and","average","gradients","across","all","GPUs",".","The","final","gradient","is","(","sum","of","gradients","across","all","GPUs",")","\/","num_frames",".","-","Update","parameters","on","each","GPU",".","Doing","so","allows","us","to","account","for","the","variability","in","number","of","masked","frames","in","variable-length","sequential","data","."]
code: def training_step(self, batch: Batch, batch_idx):
      """Custom training step with loss normalization and automatic mixed precision training.
  
      By default, DDP does the following on each train step:
      - For each GPU, compute loss and gradient on shard of training data.
      - Sync and average gradients across all GPUs. The final gradient
        is (sum of gradients across all GPUs) / N, where N is the world
        size (total number of GPUs).
      - Update parameters on each GPU.
  
      Here, we do the following:
      - For k-th GPU, compute loss and scale it by (N / num_frames), where num_frames is
        the sum of masked frames across all GPUs. Compute gradient from scaled loss.
      - Sync and average gradients across all GPUs. The final gradient
        is (sum of gradients across all GPUs) / num_frames.
      - Update parameters on each GPU.
  
      Doing so allows us to account for the variability in number of masked frames in
      variable-length sequential data.
      """
      opt = self.optimizers()
      opt.zero_grad()
      with torch.cuda.amp.autocast(enabled=True):
          loss, num_frame = self._step(batch, batch_idx, "train")
      if torch.isinf(loss) or torch.isnan(loss):
          opt.zero_grad()
          return None
  
      # 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
  
      # backward the loss and clip the gradients
      loss = self.scaler.scale(loss)
      self.manual_backward(loss)
      self.scaler.unscale_(opt)
      if self.clip_norm is not None:
          torch.nn.utils.clip_grad_norm_(
              self.model.parameters(), self.clip_norm
          )
  
      # optimization
      self.scaler.step(opt)
      sch = self.lr_schedulers()
      sch.step()
      self.scaler.update()
      return loss
  
imports: import math
  from typing import Optional, Tuple
  import torch
  import torch.nn.functional
  import torchaudio
  import torchaudio.models.wav2vec2.components
  from dataset import _get_lengths_librilightlimited, _get_lengths_librispeech, BucketizeBatchSampler, CollateFnHubert, CollateFnLibriLightLimited, DistributedBatchSampler, HuBERTDataSet
  from lightning.pytorch import LightningModule
  from loss import hubert_loss
  from torch import Tensor
  from torch.optim.optimizer import Optimizer
  from torch.utils.data import DataLoader
label_desc: Use image node_id 4 for calling the HuBERTPreTrainModule obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss, None