content: #!/usr/bin/env python3
  # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  
  """Loss functions."""
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  
  class SoftTargetCrossEntropy(nn.Module):
      """
      Cross entropy loss with soft target.
      """
  
      def __init__(self, reduction="mean"):
          """
          Args:
              reduction (str): specifies reduction to apply to the output. It can be
                  "mean" (default) or "none".
          """
          super(SoftTargetCrossEntropy, self).__init__()
          self.reduction = reduction
  
      def forward(self, x, y):
          loss = torch.sum(-y * F.log_softmax(x, dim=-1), dim=-1)
          if self.reduction == "mean":
              return loss.mean()
          elif self.reduction == "none":
              return loss
          else:
              raise NotImplementedError
  
  
  _LOSSES = {
      "cross_entropy": nn.CrossEntropyLoss,
      "bce": nn.BCELoss,
      "bce_logit": nn.BCEWithLogitsLoss,
      "soft_cross_entropy": SoftTargetCrossEntropy,
  }
  
  
  def get_loss_func(loss_name):
      """
      Retrieve the loss given the loss name.
      Args (int):
          loss_name: the name of the loss to use.
      """
      if loss_name not in _LOSSES.keys():
          raise NotImplementedError("Loss {} is not supported".format(loss_name))
      return _LOSSES[loss_name]
  