Say in the forward() pass, I generate the network predictions like the following
def forward(self, input):
    conditions = [1, 2, 3, 4, 5]
    
    out_dict = {}
    for c in conditions:
        output = self.network( (input, c) ) 
        out_dict[index of c] = output

    return out_dict

Now from the ground truth, we know that we only care about condition 3, and only want to backpropagate the loss in this condition. So I did
loss = loss_fn(out_dict[3])
loss.backward()

To verify this, I add a hook to the output of each condition. The code is shown below. If everything works perfectly, we should see “3 [gradient tensor …]”, but it is always “5 [gradient tensor …]”. So it seems like the conditional backpropagation is not working, what is going on here? and how should I do the backprop only with the tensor that I am interested in?
def forward(self, input):
    conditions = [1, 2, 3, 4, 5]
    
    out_dict = {}
    for c in conditions:
        output = self.network( (input, c) ) 
       
        output.register_hook(lambda grad: print("{index of c}", grad))

        out_dict[index of c] = output

    return out_dict