Inspired by https://arxiv.org/pdf/1711.08946.pdf 4 and https://s3-us-west-2.amazonaws.com/openai-assets/dota_benchmark_results/network_diagram_08_06_2018.pdf 2 I’m writing a reinforcement learning program and I’m implementing multiple action heads sharing a common network.

I want to propagate loss1, loss2 and loss3 to the output of the shared network and MSE(loss1,loss2,loss3) from the output of the shared network all the way back. The shared network and each of the action head networks have their own network classes and torch.optim.
I’ve tried
optim_head1.zero_grad(retain_graph=False)
loss1 = loss_fn(x1,y1)
loss1.backward()
optim_head1.step()

:

optim_head2.zero_grad(retain_graph=False)
loss2 = loss_fn(x2,y2)
loss2.backward()
optim_head2.step()

:

optim_head3.zero_grad(retain_graph=False)
loss3 = loss_fn(x3,y3)
loss3.backward()
optim_head3.step()

But I’m getting this error during the second backward() call

“Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time.”

How can I do this? What am I not understanding about how Pytorch performs backprop? Thanks.