I’m having trouble figuring out how to implement something I want in PyTorch: path-conditional gradient backpropagation.
For simplicity, suppose I have data with shape (batch size, input_dimension) and I have a simple network that outputs a scalar sum of two affine transformations of the input i.e.
linear1 = nn.Linear(in_features=input_dimension, out_features=1)
linear2 = nn.Linear(in_features=input_dimension, out_features=1)
y = linear1(x) + linear2(x)
loss = torch.mean((y - y_target) ** 2)

During backprop, I’d like to update linear1's parameters using only elements in the batch where $y < 0$ and update linear2's parameters using only elements in the batch where $y > 0$.
How can I implement this?
I’ve tried register_backward_hook, but if I’m understand the functionality correctly, by the time the registered function is called, the gradient of the error with respect to the parameters has already been calculated. I tried register_hook, but this doesn’t permit me to conditionally mask the gradient dL/dy depending on which linear layer has backward() called next.