In PyTorch, when using loss.backward() to compute gradients, how can I prevent it from overriding the gradients I've manually computed?

Question

I'm currently working on an RNN neural network for a speech recognition task where I've designed an algorithm to calculate the gradients for w_in and w_rec myself. However, I want to let PyTorch's automatic differentiation handle the gradient computation for w_out by using loss.backward(). I would like to know how to make the loss.backward() call compute the gradient only for a specific weight without affecting all parameters, thereby preventing my manually calculated gradients from being overwritten.

Add a flag variable in the model class to indicate whether gradients should be accumulated. Initialize it as False.

In the forward function, decide whether to manually calculate the gradients for w_in and w_rec based on the value of accumulate_gradients. If accumulate_gradients is True, manually calculate the gradients; otherwise, do not perform manual calculations.

In the training loop, before calling loss.backward(), set accumulate_gradients to True.

After updating the parameters in the training loop, reset accumulate_gradients to False to prevent gradients from being recalculated in the next iteration.

Will it work？