docstr: Custom training step.
  
  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 / B_total), where B_total is
    the sum of batch sizes 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) / B_total.
  - Update parameters on each GPU.
  
  Doing so allows us to account for the variability in batch sizes that
  variable-length sequential data commonly yields.
docstr_tok: ["Custom","training","step",".","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","\/","B_total",")",",","where","B_total","is","the","sum","of","batch","sizes","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",")","\/","B_total",".","-","Update","parameters","on","each","GPU",".","Doing","so","allows","us","to","account","for","the","variability","in","batch","sizes","that","variable-length","sequential","data","commonly","yields","."]
code: def training_step(self, batch: Batch, batch_idx):
      """Custom training step.
  
      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 / B_total), where B_total is
        the sum of batch sizes 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) / B_total.
      - Update parameters on each GPU.
  
      Doing so allows us to account for the variability in batch sizes that
      variable-length sequential data commonly yields.
      """
  
      opt = self.optimizers()
      opt.zero_grad()
      loss = self._step(batch, batch_idx, "train")
      batch_size = batch.features.size(0)
      batch_sizes = self.all_gather(batch_size)
      self.log(
          "Gathered batch size",
          batch_sizes.sum(),
          on_step=True,
          on_epoch=True,
          batch_size=batch.targets.size(0),
      )
      loss *= (
          batch_sizes.size(0) / batch_sizes.sum()
      )  # world size / batch size
      self.manual_backward(loss)
      torch.nn.utils.clip_grad_norm_(self.model.parameters(), 10.0)
      opt.step()
  
      # step every epoch
      sch = self.lr_schedulers()
      if self.trainer.is_last_batch:
          sch.step()
  
      return loss
  
imports: import logging
  import math
  from collections import namedtuple
  from typing import List, Tuple
  import sentencepiece
  import torch
  import torchaudio
  from pytorch_lightning import LightningModule
  from torchaudio.prototype.models import conformer_rnnt_biasing_base, Hypothesis, RNNTBeamSearchBiasing
label_desc: Use image node_id 5 for calling the ConformerRNNTModule obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss