text: from torch.nn.modules.loss import _Loss
  
  __all__ = ["JointLoss", "WeightedLoss"]
  
  
  class WeightedLoss(_Loss):
      """Wrapper class around loss function that applies weighted with fixed factor.
      This class helps to balance multiple losses if they have different scales
      """
  
      def __init__(self, loss, weight=1.0):
          super().__init__()
          self.loss = loss
          self.weight = weight
  
      def forward(self, *input):
          return self.loss(*input) * self.weight
  
  
  class JointLoss(_Loss):
      def __init__(self, first, second, first_weight=1.0, second_weight=1.0):
          super().__init__()
          self.first = WeightedLoss(first, first_weight)
          self.second = WeightedLoss(second, second_weight)
  
      def forward(self, *input):
          return self.first(*input) + self.second(*input)
  