class_docstr: Knowledge distillation loops for training and evaluation. This module is independent of
  the model definition. This class inherits from the Pytorch Lightning's LightningModule:
  https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
code: def on_before_optimizer_step(self, optimizer):
      # If using mixed precision, the gradients are already unscaled here
      if self.track_grad_norm != -1:
          self.log_dict(grad_norm(self, norm_type=self.track_grad_norm))
  
imports: import logging
  from typing import Callable, List, Optional, Union
  import lightning.pytorch
  import torch
  import torch.nn.functional
  import torchmetrics
  from lightning.pytorch.utilities import grad_norm
  from omegaconf import DictConfig
  from torch import nn
  from torch.nn.modules.loss import _Loss
  from torchmetrics.aggregation import BaseAggregator
  from ..constants import FEATURES, LOGITS, WEIGHT
  from ..models.utils import run_model
  from .utils import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler, get_optimizer
label_desc: Use image node_id 13 for calling the DistillerLitModule obj's underlying member method code with example usage: obj.on_before_optimizer_step(optimizer) without return types