You can calculate the loss, and before doing the reduction over the batch, multiply it with the scaling factor, and then do the reduction. Here’s something implemented:

  

      github.com
  

  
    dpressel/mead-baseline/blob/master/layers/eight_mile/pytorch/layers.py#L4693



      
      
              """
              def __init__(self):
                  super().__init__()
                  self.loss = nn.NLLLoss(reduction='none')
          
          
    def forward(self, preds, targets, weights):
                  loss = sum([self.loss(pred, targets[:, i]) for i, pred in enumerate(preds)])
                  weights = weights.type_as(loss)
                  return torch.dot(loss, weights)/len(weights)
          
          
class WeightedSequenceLoss(nn.Module):
              """Weight individual training examples
          
          
    """
              def __init__(self, LossFn: nn.Module = nn.NLLLoss, avg: str = "token"):
                  super().__init__()
                  self.avg = avg
                  self.crit = LossFn(ignore_index=Offsets.PAD, reduction="none")
                  if avg == 'token':
                      self._reduce = self._mean
                  else:
      
    
  


  

  
    
    
  

  


This is for sequence though, you probably have BxC and not BxTxC, but I think the idea is similar.