code: def training_step(self, batch, batch_idx) -> Dict:
      loss_tensors = self._step(batch)
  
      logs = {
          name: loss
          for name, loss in zip(self.loss_names, loss_tensors)
      }
      # tokens per batch
      logs["tpb"] = (
          batch["input_ids"].ne(self.pad).sum()
          + batch["labels"].ne(self.pad).sum()
      )
      logs["bs"] = batch["input_ids"].shape[0]
      logs["src_pad_tok"] = batch["input_ids"].eq(self.pad).sum()
      logs["src_pad_frac"] = (
          batch["input_ids"].eq(self.pad).float().mean()
      )
      # TODO(SS): make a wandb summary metric for this
      return {"loss": loss_tensors[0], "log": logs}
  
imports: import argparse
  import glob
  import logging
  import os
  import sys
  import time
  from collections import defaultdict
  from pathlib import Path
  from typing import Dict, List, Tuple
  import numpy
  import pytorch_lightning
  import torch
  from torch import nn
  from torch.utils.data import DataLoader
  from callbacks import Seq2SeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
  from transformers import MBartTokenizer, T5ForConditionalGeneration
  from transformers.models.bart.modeling_bart import shift_tokens_right
  from utils import ROUGE_KEYS, LegacySeq2SeqDataset, Seq2SeqDataset, assert_all_frozen, calculate_bleu, calculate_rouge, check_output_dir, flatten_list, freeze_embeds, freeze_params, get_git_info, label_smoothed_nll_loss, lmap, pickle_save, save_git_info, save_json, use_task_specific_params
  from lightning_base import BaseTransformer, add_generic_args, generic_train
label_desc: Use image node_id 7 for calling the SummarizationModule obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: dict