I have a use case in which I have multiple loss functions on which I have to call backward without using any reduction like mean or sum. I want to calculate gradients for different losses parallely.
losses = [loss1, loss2, loss3] losses.backward() print(param.grad) ## It should contain the jacobian of the gradients
In particular I am trying to implement this functionality in pytorch https://github.com/tianheyu927/PCGrad/blob/c5fbd7c856526373828074f06875230f7f3ee79e/PCGrad_tf.py#L39
Is it possible to do it parallely without a for loop?