For efficiency, you’ll want to use tensor operations instead of looping over every element.
Something like:
output = ... # batch x 2 x 180 x 320
u = output[:, 0, :, :]
v = output[:, 1, :, :]

def delta_neighbors(x):
  assert x.dim() == 3
  x = x.unsqueeze(1)  # view as 4D: batch x 1 x h x w
  padded = F.pad(x, (1, 1, 1, 1))  # pad with zeros on each side
  padded = padded.squeeze(1)  # view as 3D
  top = padded[:, :-2, 1:-1]
  bottom = padded[:, 2:, 1:-1]
  left = padded[:, 1:-1, :-2]
  right = padded[:, 1:-1, 2:]
  return (padded - top).abs() * c[0] + (padded - bottom).abs() * c[1] + (padded - left).abs() * c[2] + (padded - right).abs() * c[3]

loss = delta_neighbors(u).sum() + delta_neighbors(v).sum()