Naruto-Sasuke:

def SoftCrossEntropy(inputs, target, reduction='sum'):
    log_likelihood = -F.log_softmax(inputs, dim=1)
    batch = inputs.shape[0]
    if reduction == 'average':
        loss = torch.sum(torch.mul(log_likelihood, target)) / batch
    else:
        loss = torch.sum(torch.mul(log_likelihood, target))
    return loss



Thanks. I have tried to set the labels in the loss forward method, but this did not work - any idea how to do it please? My code is below:
class SoftCrossEntropy(nn.Module):
    def __init__(self, labels, reduction='sum'):
        super().__init__()
        self.labels = labels
        self.reduction = reduction
    
    def forward(self, preds, target):
        target = labels[target]
        log_likelihood = -F.log_softmax(preds, dim=1)
        batchsize = preds.shape[0]
        if self.reduction == 'mean':
            loss = tc.sum(tc.mul(log_likelihood, target)) / batchsize
        else:
            loss = tc.sum(tc.mul(log_likelihood, target))
        return loss

class ResNetMNIST_ls(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = models.resnet18(num_classes=10)
        self.model.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2),
                                     padding=(3, 3), bias=False)
        self.loss = SoftCrossEntropy(labels)

    @auto_move_data
    def forward(self, x):
        return self.model(x)
  
    def training_step(self, batch, batch_no):
        x, y = batch
        logits = self(x)
        loss = self.loss(logits, y)
        return loss
    
    def validation_step(self, batch, batch_no):
        x, y = batch
        logits = self(x)
        loss = self.loss(logits, y)
  
    def configure_optimizers(self):
        return tc.optim.Adam(self.parameters(), lr=1e-3)