Hmmm… I’m not sure I understand especially if I’m trying to use gradients calculated from the loss of two different tasks. Another way to ask what I’m trying to do is: what is the pytorch equivalent of this code?
#Where x represents an input to a model

 with tf.GradientTape(persistent=True) as tape:
      primary_loss = compute_primary_task_loss(x)
      auxiliary_loss = compute_auxiliary_task_loss(x)
    
    primary_grad = tape.gradient(primary_loss, x)
    auxiliary_grad = tape.gradient(auxiliary_loss, x)
    new_grad = censored_vector(auxiliary_grad,
                               primary_grad,
                               mode=mode)
    optimizer.apply_gradients([(primary_grad + lam*new_grad, x)])