Hi, I noticed that the output of cross-entropy loss (for semantic segmentation use case so K-dimensional one) with reduction="mean" is different than when I calculate it with sum and mean on unreduced output. This is most visible with a bigger batch size.
A minimal working example:
import torch
import torch.nn as nn
import numpy as np

basic_img = torch.Tensor([arr for arr in np.random.rand(128, 2, 768, 768)])
label_basic = torch.Tensor([arr for arr in np.random.randint(2, size=(128,  768, 768))]).long()

criterion = nn.CrossEntropyLoss(reduction='none', weight=None)
loss = criterion(basic_img, label_basic)
loss.mean()

>>> tensor(0.7137)

criterion = nn.CrossEntropyLoss(reduction='mean', weight=None)
loss = criterion(basic_img, label_basic)
loss

>>> tensor(1.6724)

As far as I understand, these two versions should be equal. Am I missing something?
Pytorch version: 1.8.1
Thank you