Hi Stefano,
Sorry for the delayed reply!
It was a long time ago that I experimented with this and I haven’t pursued the issue much further but here is a workaround that helped for me, but this is for a 2 tasks example only (I wouldn’t speak of memory “leak” by the way, retain_graph=True just requires lots of memory.)
So suppose your total loss is:
loss = w_1 * loss_1 + w_2 * loss_2
# step A: calculate gradient (and its magnitude) of FIRST loss wrt a certain layer "manually", WITH retain_graph=True

gw_1 = torch.autograd.grad(
    loss_1,
    (
        myNet.<whichever layer is the last common one>.conv.weight,
        myNet.<whichever layer is the last common one>.conv.bias,
    ),
    retain_graph=True,
)

G_W_1 = torch.sqrt(torch.sum(gw_1[0] ** 2) + torch.sum(gw_1[1] ** 2))

# step B: BACKWARD through ENTIRE network
optimizer.zero_grad()
loss.backward()
optimizer.step()

# step C: now calculate gradient of SECOND loss by simply SUBTRACTING the FIRST gradient
# from the SUM that was calculated via the backward on the TOTAL loss
# (alternative would be another autograd.grad on loss_2 with retain_graph=True followed by the loss.backward,
# but that's where memory problems started for me)


gw_2 = []
gw_2.append(myNet.<whichever layer is the last common one>.conv.weight.grad - gw_1[0])
gw_2.append(myNet.<whichever layer is the last common one>.conv.bias.grad - gw_1[1])

G_W_2 = torch.sqrt(torch.sum(gw_2[0] ** 2) + torch.sum(gw_2[1] ** 2))

G_W_mean = 0.5 * (G_W_1 + G_W_2)
G_W_mean_no_grad = G_W_mean.detach()

Then you continue to work with these two quantities, G_W_mean, G_W_1 and G_W_2 to follow the loss-weight update strategy for w_1 and w_2 as described in the paper.
I guess this only works for 2 task problems though.
Anyway, I hope it helped a little.
Cheers,
Bernie