I get a discrepancy between the values of the losses obtained by the torch. CrossEntropyLoss and CustomeCrossEntropyLoss.
Issue on github 1
Code to Reproduce
import torch
import torch.nn.modules.loss as L  

class CustomCrossEntropyLoss(L._Loss):
    def __init__(self, reduction=True):
        super(CustomCrossEntropyLoss, self).__init__()
        self.reduction = reduction

     def forward(self, inp, target):
         input_target = inp.gather(1, target.view(-1, 1))
         input_max, _ = inp.max(dim=1, keepdim=True)
         output_exp = torch.exp(inp - input_max)
         output_softmax_sum = output_exp.sum(dim=1)
         output = -input_target + torch.log(output_softmax_sum).view(-1, 1) + input_max
         if self.reduction:
             output = output.mean()
         return output

 torch_ce = torch.nn.CrossEntropyLoss(reduction='none')
 custom_ce = CustomCrossEntropyLoss(reduction=False)
 batch_size = 128
 N_class = 90000
 logits = torch.randn((batch_size,N_class))
 targets = torch.randint(N_class, (batch_size,))
 print((torch_ce(logits, targets).view(-1) - custom_ce(logits,targets).view(-1)).mean())

I get a minimal non-zero discrepancy of the order of 10e-7, which then occurs in gradients and can affect the learning outcome.
How to fix this problem?