I did read many posts on the website discussing the alternatives in PyTorch to implement a custom loss function. I would like to confirm that the current state of the project works just fine with the implementation below:
class TVLoss(Module):
    """
    Total variation loss.
    """
    def __init__(self):
        super(TVLoss, self).__init__()

    def forward(self, yhat, y):
        bsize, chan, height, width = y.size()
        errors = []
        for h in range(height-1):
            dy = torch.abs(y[:,:,h+1,:] - y[:,:,h,:])
            dyhat = torch.abs(yhat[:,:,h+1,:] - yhat[:,:,h,:])
            error = torch.norm(dy - dyhat, 1)
            errors.append(error)

        return sum(errors) / height

Is there any inefficiency or missing component? The training is apparently working but I didn’t have the chance to train for long enough.