class_docstr: Compute the augmentation loss, as described in the
  `RangeAugment <https://arxiv.org/abs/2212.10553>`_ paper.
  
  Args:
      opts: command line arguments
class_docstr_tok: ["Compute","the","augmentation","loss",",","as","described","in","the","`","RangeAugment","<","https",":","\/\/arxiv.org\/abs\/2212.10553",">","`","_","paper",".","Args",":","opts",":","command","line","arguments"]
docstr: Compute the loss between input and augmented image, as described in
  `RangeAugment <https://arxiv.org/abs/2212.10553>`_ paper.
  
  Args:
      input_sample: Input sample can either be a Tensor or a dictionary with mandatory key "image". In
          case of a dictionary, the values can be a Tensor or list of Tensors.
      prediction: Output of augmentation model. Mapping of (string: Tensor) with `augmented_tensor`
          as the required key.
  
  Shapes:
      input_sample:
          * Tensor: The shape of input tensor is [N, C, H, W]
          * Mapping[str, Tensor]: The shape of tensor is [N, C, H, W]
          * Mapping[str, List[Tensor]]: The length of List is N, and the shape of each tensor is [1, C, H, W]
      prediction: The shape of prediction["augmented_tensor"] is [N, C, H, W]
  
  Returns:
      A scalar loss value
  
  ...note:
      During validation or evaluation, neural augmentation loss is not computed and 0 is returned
docstr_tok: ["Compute","the","loss","between","input","and","augmented","image",",","as","described","in","`","RangeAugment","<","https",":","\/\/arxiv.org\/abs\/2212.10553",">","`","_","paper",".","Args",":","input_sample",":","Input","sample","can","either","be","a","Tensor","or","a","dictionary","with","mandatory","key","``","image","''",".","In","case","of","a","dictionary",",","the","values","can","be","a","Tensor","or","list","of","Tensors",".","prediction",":","Output","of","augmentation","model",".","Mapping","of","(","string",":","Tensor",")","with","`","augmented_tensor","`","as","the","required","key",".","Shapes",":","input_sample",":","*","Tensor",":","The","shape","of","input","tensor","is","[","N",",","C",",","H",",","W","]","*","Mapping","[","str",",","Tensor","]",":","The","shape","of","tensor","is","[","N",",","C",",","H",",","W","]","*","Mapping","[","str",",","List","[","Tensor","]","]",":","The","length","of","List","is","N",",","and","the","shape","of","each","tensor","is","[","1",",","C",",","H",",","W","]","prediction",":","The","shape","of","prediction","[","``","augmented_tensor","''","]","is","[","N",",","C",",","H",",","W","]","Returns",":","A","scalar","loss","value","...","note",":","During","validation","or","evaluation",",","neural","augmentation","loss","is","not","computed","and","0","is","returned"]
code: def forward(
      self,
      input_sample: Union[
          Tensor, Mapping[str, Union[Tensor, List[Tensor]]]
      ],
      prediction: Mapping[str, Tensor],
      *args,
      **kwargs,
  ) -> Tensor:
      """Compute the loss between input and augmented image, as described in
      `RangeAugment <https://arxiv.org/abs/2212.10553>`_ paper.
  
      Args:
          input_sample: Input sample can either be a Tensor or a dictionary with mandatory key "image". In
              case of a dictionary, the values can be a Tensor or list of Tensors.
          prediction: Output of augmentation model. Mapping of (string: Tensor) with `augmented_tensor`
              as the required key.
  
      Shapes:
          input_sample:
              * Tensor: The shape of input tensor is [N, C, H, W]
              * Mapping[str, Tensor]: The shape of tensor is [N, C, H, W]
              * Mapping[str, List[Tensor]]: The length of List is N, and the shape of each tensor is [1, C, H, W]
          prediction: The shape of prediction["augmented_tensor"] is [N, C, H, W]
  
      Returns:
          A scalar loss value
  
      ...note:
          During validation or evaluation, neural augmentation loss is not computed and 0 is returned
      """
      if not self.training:
          return torch.tensor(
              0.0, device=self.device, dtype=torch.float
          )
  
      if not isinstance(prediction, Mapping):
          logger.error(
              "Prediction needs to be an instance of Mapping and must contain augmented_tensor"
              " as keys"
          )
      if isinstance(input_sample, Mapping):
          input_sample = input_sample["image"]
  
      if isinstance(input_sample, List):
          # if its a list of images, stack them
          input_sample = torch.stack(input_sample, dim=0)
  
      augmented_tensor = prediction["augmented_tensor"]
      loss_na = self._compute_loss(
          input_tensor=input_sample,
          augmented_tensor=augmented_tensor,
          *args,
          **kwargs,
      )
  
      return loss_na
  
file_path: ml-cvnets/loss_fn/neural_augmentation.py
imports: import argparse
  import math
  from typing import List, Mapping, Union
  import torch
  from torch import Tensor
  from torch.nn import functional
  from loss_fn import LOSS_REGISTRY, BaseCriteria
  from utils import logger
  from utils.ddp_utils import is_master
label_desc: Use image node_id 5 for calling the NeuralAugmentation obj's underlying member method code with example usage: obj.forward(input_sample, prediction) and returns: loss_na, torch