class_docstr: Base class for defining distillation loss functions. Sub-classes must implement `_forward_distill` function.
  
  Args:
      opts: command line arguments
docstr: Computes distillation loss
  
  Args:
      input_sample: Input image tensor.
      prediction: Output of model. It can be a tensor or mapping of (string: Tensor). In case of a dictionary,
      `logits` is a required key.
      target: Target label tensor containing values in the range `[0, C)`, where :math:`C` is the number of classes
  
  Shapes:
      input_sample: The shape of input tensor is [N, C, H, W]
      prediction:
          * When prediction is a tensor, then shape is [N, C]
          * When prediction is a dictionary, then shape of prediction["logits"] is [N, C]
      target: The shape of target tensor is [N]
  
  Returns:
      * Scalar loss value is returned.
docstr_tok: ["Computes","distillation","loss","Args",":","input_sample",":","Input","image","tensor",".","prediction",":","Output","of","model",".","It","can","be","a","tensor","or","mapping","of","(","string",":","Tensor",")",".","In","case","of","a","dictionary",",","`","logits","`","is","a","required","key",".","target",":","Target","label","tensor","containing","values","in","the","range","`","[","0",",","C",")","`",",","where",":","math",":","`","C","`","is","the","number","of","classes","Shapes",":","input_sample",":","The","shape","of","input","tensor","is","[","N",",","C",",","H",",","W","]","prediction",":","*","When","prediction","is","a","tensor",",","then","shape","is","[","N",",","C","]","*","When","prediction","is","a","dictionary",",","then","shape","of","prediction","[","``","logits","''","]","is","[","N",",","C","]","target",":","The","shape","of","target","tensor","is","[","N","]","Returns",":","*","Scalar","loss","value","is","returned","."]
code: def forward(
      self,
      input_sample: Tensor,
      prediction: Union[Mapping[str, Tensor], Tensor],
      target: Tensor,
      *args,
      **kwargs,
  ) -> Union[Mapping[str, Tensor], Tensor]:
      """Computes distillation loss
  
      Args:
          input_sample: Input image tensor.
          prediction: Output of model. It can be a tensor or mapping of (string: Tensor). In case of a dictionary,
          `logits` is a required key.
          target: Target label tensor containing values in the range `[0, C)`, where :math:`C` is the number of classes
  
      Shapes:
          input_sample: The shape of input tensor is [N, C, H, W]
          prediction:
              * When prediction is a tensor, then shape is [N, C]
              * When prediction is a dictionary, then shape of prediction["logits"] is [N, C]
          target: The shape of target tensor is [N]
  
      Returns:
          * Scalar loss value is returned.
      """
  
      if isinstance(prediction, Tensor):
          return self._forward_distill(
              input_sample=input_sample,
              prediction=prediction,
              *args,
              **kwargs,
          )
      elif isinstance(prediction, Mapping):
          if "logits" not in prediction:
              logger.error(
                  f"logits is a required key in {self.__class__.__name__} when prediction type"
                  f"is dictionary. Got keys: {prediction.keys()}"
              )
  
          predicted_logits = prediction["logits"]
  
          # compute distillation loss
          distill_loss = self._forward_distill(
              input_sample=input_sample,
              prediction=predicted_logits,
              *args,
              **kwargs,
          )
          return distill_loss
      else:
          logger.error(
              f"Prediction should be either a Tensor or Dictionary[str, Tensor]. Got: {type(prediction)}"
          )
  
file_path: ml-cvnets/loss_fn/distillation/base_distillation.py
imports: import argparse
  from typing import Mapping, Union
  import torch
  from torch import Tensor
  from loss_fn import LOSS_REGISTRY, BaseCriteria
  from loss_fn.utils.build_helper import build_cls_teacher_from_opts
  from utils import logger
label_desc: Use image node_id 5 for calling the BaseDistillationCriteria obj's underlying member method code with example usage: obj.forward(input_sample, prediction, target) and returns: self, distill_loss