For 2, you can implement the vector of squared distances using PyTorch mathematical ops, and then replace the elements of prediction with the corresponding elements of target for the indices that you don’t want to backpropagate. (It’ll still backprop them, but with gradient zero).