class_docstr: This class defines a custom loss function for product of predictions and labels.
  
  Attributes:
  alpha: A float indicating how much to weigh the weak model.
  beta: A float indicating how much to weigh the strong model.
  warmup_frac: A float indicating the fraction of total training steps for warmup.
class_docstr_tok: ["This","class","defines","a","custom","loss","function","for","product","of","predictions","and","labels",".","Attributes",":","alpha",":","A","float","indicating","how","much","to","weigh","the","weak","model",".","beta",":","A","float","indicating","how","much","to","weigh","the","strong","model",".","warmup_frac",":","A","float","indicating","the","fraction","of","total","training","steps","for","warmup","."]
code: def __call__(
      self,
      logits: torch.Tensor,
      labels: torch.Tensor,
      step_frac: float,
  ) -> torch.Tensor:
      preds = torch.softmax(logits, dim=-1)
      target = torch.pow(preds, self.beta) * torch.pow(
          labels, self.alpha
      )
      target /= target.sum(dim=-1, keepdim=True)
      target = target.detach()
      loss = torch.nn.functional.cross_entropy(
          logits, target, reduction="none"
      )
      return loss.mean()
  
file_path: weak-to-strong/weak_to_strong/loss.py
imports: import torch
label_desc: Use image node_id 2 for calling the product_loss_fn obj's underlying member method code with example usage: obj.__call__(logits, labels, step_frac) and returns: loss