class_docstr: Base class to calculate hidden state loss.
  Returned dict includes a loss value and additional outputs.
class_docstr_tok: ["Base","class","to","calculate","hidden","state","loss",".","Returned","dict","includes","a","loss","value","and","additional","outputs","."]
docstr: A wrapper around custom _loss that adds a weighted loss and name to the output dict
docstr_tok: ["A","wrapper","around","custom","_loss","that","adds","a","weighted","loss","and","name","to","the","output","dict"]
code: def loss(self, inputs, batch_data=None):
      """A wrapper around custom _loss that adds a weighted loss and name to the output dict"""
      self._validate_inputs(inputs)
  
      loss_dict = self._loss(inputs, batch_data=batch_data)
      if "loss" not in loss_dict:
          raise KeyError("Loss dict must contain 'loss' key")
  
      # average loss over active steps only. loss [B x S]
      loss = loss_dict["loss"]
      # hiddens_mask has shape of [B x S]
      hiddens_mask = inputs["hiddens_mask"].to(loss)
      loss = loss * hiddens_mask
      # sequence level loss [B x S] -> batch level loss [B]
      loss = loss.sum(dim=1) / hiddens_mask.sum(dim=1).clamp(min=1.0)
  
      # compute batch level weighted loss (scalar)
      weighted_loss = loss.sum() * self.loss_weight
  
      # store updated losses
      loss_dict["loss"] = loss
      loss_dict["weighted_loss"] = weighted_loss
      loss_dict["weight_loss"] = torch.tensor(self.loss_weight).to(
          weighted_loss
      )
  
      return loss_dict
  
imports: import math
  import torch
label_desc: Use image node_id 7 for calling the MegatronBaseHiddenLoss obj's underlying member method code with example usage: obj.loss(inputs, batch_data) and returns: loss_dict