In my training loop what would the affects of truncating the model output and target vectors as so:
action_output = model(inputs)
         
# The effects of the below two lines
action_output = action_output[:, tasks]   # Tasks is an array of 1*"target size"
action_target = action_target[:, tasks]

action_loss = criterion(action_output, action_target)
        
action_loss.backward()

What I believe it is doing is that it is truncating the computation graph and backproping only on weights that affect that output node.