Solved by colesbury in post #2 
              
              
                This looks fine. If your height is large, you may get better performance if you avoid the for-loop: 
def forward(self, yhat, y):
    bsize, chan, height, width = y.size()
    errors = []
    dy = torch.abs(y[:,:,1:,:] - y[:,:,:-1,:])
    dyhat = torch.abs(yhat[:,:,1:,:] - yhat[:,:,:-1,:])
    error …