class_docstr: Cross entropy loss for the task of semantic segmentation.
  
  Args:
      opts: command-line arguments
docstr: Computes the segmentation loss. If prediction is a Tuple[Tensor, Tensor], then weighted sum of CE losses is
  computed.
  
  Args:
      prediction: Output of segmentation model. If auxiliary branch is enabled, then prediction is
      a Tuple[Tensor, Tensor]. Otherwise, it is a Tensor.
      target: Ground truth segmentation mask.
  
  Shapes:
      prediction:
          * When prediction is a Tensor, then shape is [Batch size, Channels, Height, Width]
          * When prediction is a Tuple[Tensor, Tensor], then shape of one tensor is [Batch size, Channels, Height, Width]
              while the other is [Batch size, Channels, Height / O, Width/ O]
              where O is output stride of feature map (typically 4).
      target: Shape is [Batch size, Height, Width]
  
  Returns:
      Mapping of the form (string: scalar value) is returned with total_loss as mandatory and
      (seg_loss, aux_loss) as optional keys. total_loss is weighted sum of seg_loss and aux_loss (when applicable).
  
  ...note:
      When shape of prediction and target are not the same, prediction is resized using bilinear interpolation to
      match the size of target.
docstr_tok: ["Computes","the","segmentation","loss",".","If","prediction","is","a","Tuple","[","Tensor",",","Tensor","]",",","then","weighted","sum","of","CE","losses","is","computed",".","Args",":","prediction",":","Output","of","segmentation","model",".","If","auxiliary","branch","is","enabled",",","then","prediction","is","a","Tuple","[","Tensor",",","Tensor","]",".","Otherwise",",","it","is","a","Tensor",".","target",":","Ground","truth","segmentation","mask",".","Shapes",":","prediction",":","*","When","prediction","is","a","Tensor",",","then","shape","is","[","Batch","size",",","Channels",",","Height",",","Width","]","*","When","prediction","is","a","Tuple","[","Tensor",",","Tensor","]",",","then","shape","of","one","tensor","is","[","Batch","size",",","Channels",",","Height",",","Width","]","while","the","other","is","[","Batch","size",",","Channels",",","Height","\/","O",",","Width\/","O","]","where","O","is","output","stride","of","feature","map","(","typically","4",")",".","target",":","Shape","is","[","Batch","size",",","Height",",","Width","]","Returns",":","Mapping","of","the","form","(","string",":","scalar","value",")","is","returned","with","total_loss","as","mandatory","and","(","seg_loss",",","aux_loss",")","as","optional","keys",".","total_loss","is","weighted","sum","of","seg_loss","and","aux_loss","(","when","applicable",")",".","...","note",":","When","shape","of","prediction","and","target","are","not","the","same",",","prediction","is","resized","using","bilinear","interpolation","to","match","the","size","of","target","."]
code: def _forward_seg(
      self,
      prediction: Union[Tensor, Tuple[Tensor, Tensor]],
      target: Tensor,
      *args,
      **kwargs,
  ) -> Mapping[str, Tensor]:
      """Computes the segmentation loss. If prediction is a Tuple[Tensor, Tensor], then weighted sum of CE losses is
      computed.
  
      Args:
          prediction: Output of segmentation model. If auxiliary branch is enabled, then prediction is
          a Tuple[Tensor, Tensor]. Otherwise, it is a Tensor.
          target: Ground truth segmentation mask.
  
      Shapes:
          prediction:
              * When prediction is a Tensor, then shape is [Batch size, Channels, Height, Width]
              * When prediction is a Tuple[Tensor, Tensor], then shape of one tensor is [Batch size, Channels, Height, Width]
                  while the other is [Batch size, Channels, Height / O, Width/ O]
                  where O is output stride of feature map (typically 4).
          target: Shape is [Batch size, Height, Width]
  
      Returns:
          Mapping of the form (string: scalar value) is returned with total_loss as mandatory and
          (seg_loss, aux_loss) as optional keys. total_loss is weighted sum of seg_loss and aux_loss (when applicable).
  
      ...note:
          When shape of prediction and target are not the same, prediction is resized using bilinear interpolation to
          match the size of target.
      """
  
      aux_out = None
      if isinstance(prediction, Tuple) and len(prediction) == 2:
          mask, aux_out = prediction
          assert isinstance(mask, Tensor)
          assert isinstance(aux_out, Tensor)
      elif isinstance(prediction, Tensor):
          mask = prediction
          assert isinstance(mask, Tensor)
      else:
          raise NotImplementedError(
              "For computing loss for segmentation task, we need prediction to be an instance of Tuple or Tensor"
          )
  
      cls_wts = None
      if self.training:
          if self.use_class_wts:
              n_classes = mask.size(1)  # Mask is of shape B x C x H x W
              cls_wts = compute_class_weights(
                  target=target, n_classes=n_classes
              )
          seg_loss = self._compute_loss(
              pred_mask=mask, target_mask=target, weight=cls_wts
          )
  
          if aux_out is not None:
              loss_aux = self._compute_loss(
                  pred_mask=aux_out, target_mask=target, weight=cls_wts
              )
              total_loss = seg_loss + (self.aux_wt * loss_aux)
              return {
                  "total_loss": total_loss,
                  "seg_loss": seg_loss,
                  "aux_loss": (self.aux_wt * loss_aux),
              }
          return {"total_loss": seg_loss}
      else:
          # during validation, we do not compute aux. loss
          seg_loss = self._compute_loss(
              pred_mask=mask, target_mask=target, weight=None
          )
          return {"total_loss": seg_loss}
  
file_path: ml-cvnets/loss_fn/segmentation/cross_entropy.py
imports: import argparse
  from typing import Any, Mapping, Optional, Tuple, Union
  from torch import Tensor
  from torch.nn import functional
  from loss_fn import LOSS_REGISTRY
  from loss_fn.segmentation.base_segmentation_criteria import BaseSegmentationCriteria
  from loss_fn.utils.class_weighting import compute_class_weights
  from utils import logger
label_desc: Use image node_id 4 for calling the SegCrossEntropy obj's underlying member method code with example usage: obj._forward_seg(prediction, target) and returns: dict, dict, dict