from torch.nn.modules.loss import _WeightedLoss
import torch.nn.functional as F

class SmoothCrossEntropyLoss(_WeightedLoss):
    def __init__(self, weight=None, reduction='mean'):
        super().__init__(weight=weight, reduction=reduction)
        self.weight = weight
        self.reduction = reduction

    def reduce_loss(self, loss):
        return loss.mean() if self.reduction == 'mean' else loss.sum() \
        if self.reduction == 'sum' else loss

    def forward(self, inputs, targets):

        log_preds = F.log_softmax(inputs, -1)

        if self.weight is not None:
            log_preds = log_preds * self.weight.unsqueeze(0)

        return self.reduce_loss(-(targets * log_preds).sum(dim=-1))

maybe work. ¯\_(ツ)_/¯
class Probabilities shoud be at last dim.
import torch
import torch.nn as nn
t = torch.ones(1,1024,15)
o = torch.randn(1,1024,15)
loss_fn = SmoothCrossEntropyLoss()
loss_fn(o,t) #-> 47.3582
loss_fn1 = nn.CrossEntropyLoss()
loss_fn1(o.transpose(1,2),t.transpose(1,2)) #-> 47.3582