class_docstr: This class defines the training and validation code for training models with CVNets
code: def compute_grad_norm(self):
      parameters = [
          p for p in self.model.parameters() if p.grad is not None
      ]
      if len(parameters) == 0:
          return None
  
      norm_type = 2.0  # L2 norm
  
      inv_scale = 1.0 / self.gradient_scaler.get_scale()
      total_norm = torch.norm(
          torch.stack(
              [
                  torch.norm(p.grad.detach() * inv_scale, norm_type).to(
                      self.device
                  )
                  for p in parameters
              ]
          ),
          norm_type,
      )
      if total_norm.isnan() or total_norm.isinf():
          return None
      return total_norm
  
imports: import copy
  import gc
  import shutil
  import time
  import traceback
  from itertools import product
  from typing import Dict, Union
  import numpy
  import torch
  from torch import Tensor
  from torch.nn import functional
  from common import DEFAULT_EPOCHS, DEFAULT_ITERATIONS, DEFAULT_LOG_FREQ, if_test_env
  from data.transforms.image_torch import apply_mixing_transforms
  from engine.utils import autocast_fn, get_batch_size, get_log_writers, log_metrics
  from loss_landscape import landscape_utils
  from metrics.stats import Statistics
  from options.parse_args import parse_validation_metric_names
  from utils import logger
  from utils.checkpoint_utils import copy_weights, save_checkpoint, save_interval_checkpoint
  from utils.common_utils import move_to_device, unwrap_model_fn
  from utils.ddp_utils import dist_barrier, is_master
  from utils.tensor_utils import reduce_tensor_sum
label_desc: Use image node_id 2 for calling the Trainer obj's underlying member method code with example usage: obj.compute_grad_norm() and returns: total_norm, None, None