Hi. I’m trying to build a loss function for regression over each pixels of classes given classes and target values of pixels.
Input arguments are y_pred [N,C,H,W], classes[N,H,W], y[N,H,W].
each pixels belong to certain class which is second argument and calculate the mse loss of y_pred and y.
What I’d like to know is this function is differentiable for back propagation.
If not, Is there a way to implement? thanks in advance.
def categorical_mse_2d_loss(y_pred, classes, y):
  b, h, w = y.size()
  s = torch.zeros((b, h, w), requires_grad=True)
  for i in range(b):
    for j in range(h):
      for k in range(w):
        c = classes[i,j,k]
        s[i,j,k] = torch.square((y_pred[i,c,j,k] - y[i,j,k])) # is this differentiable? 
  
  return torch.mean(s)