class_docstr: Machine translation model which supports bottleneck architecture,
  NLL, VAE, and MIM loss.
  
  Supported losses:
    1) nll - Conditional cross entropy (the usual NMT loss)
    2) mim - MIM learning framework. A latent variable model with good
                    reconstruction and compressed latent representation.
                    https://arxiv.org/pdf/2003.02645.pdf
    3) vae - VAE learning framework. A latent variable model which learns
                    good probability estimation over observations and
                    a regularized latent representation.
                    https://arxiv.org/pdf/1312.6114.pdf
class_docstr_tok: ["Machine","translation","model","which","supports","bottleneck","architecture",",","NLL",",","VAE",",","and","MIM","loss",".","Supported","losses",":","1",")","nll","-","Conditional","cross","entropy","(","the","usual","NMT","loss",")","2",")","mim","-","MIM","learning","framework",".","A","latent","variable","model","with","good","reconstruction","and","compressed","latent","representation",".","https",":","\/\/arxiv.org\/pdf\/2003.02645.pdf","3",")","vae","-","VAE","learning","framework",".","A","latent","variable","model","which","learns","good","probability","estimation","over","observations","and","a","regularized","latent","representation",".","https",":","\/\/arxiv.org\/pdf\/1312.6114.pdf"]
docstr: Lightning calls this inside the training loop with the data from the training dataloader
  passed in as `batch`.
docstr_tok: ["Lightning","calls","this","inside","the","training","loop","with","the","data","from","the","training","dataloader","passed","in","as","`","batch","`","."]
code: def training_step(self, batch, batch_idx):
      """
      Lightning calls this inside the training loop with the data from the training dataloader
      passed in as `batch`.
      """
      # forward pass
      for i in range(len(batch)):
          if batch[i].ndim == 3:
              # Dataset returns already batched data and the first dimension of size 1 added by DataLoader
              # is excess.
              batch[i] = batch[i].squeeze(dim=0)
      src_ids, src_mask, tgt_ids, tgt_mask, labels = batch
      z, z_mean, z_logv, z_mask, tgt_log_probs = self(
          src_ids, src_mask, tgt_ids, tgt_mask
      )
      train_loss, info_dict = self.loss(
          z=z,
          z_mean=z_mean,
          z_logv=z_logv,
          z_mask=z_mask,
          tgt_log_probs=tgt_log_probs,
          tgt=tgt_ids,
          tgt_mask=tgt_mask,
          tgt_labels=labels,
          train=True,
          return_info=True,
      )
      tensorboard_logs = {
          "train_loss": train_loss,
          "lr": self._optimizer.param_groups[0]["lr"],
      }
      tensorboard_logs.update(info_dict)
  
      return {"loss": train_loss, "log": tensorboard_logs}
  
imports: from typing import Dict, Optional
  import numpy
  import torch
  from pytorch_lightning import Trainer
  from nemo.collections.common.losses import NLLLoss
  from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import MTBottleneckModelConfig
  from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
  from nemo.core.classes.common import typecheck
  from nemo.utils import timers
label_desc: Use image node_id 9 for calling the MTBottleneckModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: dict