Как в Pytorch реализовать пользовательскую функцию потерь?

Вот код функции потерь(log_loss - это потеря из sklearn), у меня проблема с backward, не могу понять как его реализовать(уже пытался это сделать через autograd function).

from torch.autograd.function import Function
class custom_loss:#(Function):
    def __init__(self, preds, targets):
        self.preds = preds
        self.targets = targets
    #@staticmethod
    def forward(self):
        log_loss_ = 0
        #self.preds.save_for_backward(self.targets)
        #gt = np.array(gt)
    
        for i in range(10):
            log_loss_ += log_loss(self.target[:, i], self.preds[:, i])
        
        return torch.tensor(log_loss_ / 10, requires_grad = True)