Thank you for this nice workaround. However to get b.grad = b1.grad + b2.grad I managed to do a hacky way as below:
b_grad = 0
for param in total_loss.parameters():
    if param.grad is not None:
      b_grad += param.grad
  
  for param in total_loss.parameters():
    if param.grad is not None:
      param.grad = b_grad

Is there a better way to do this? Because, when you make copies of the original parameter (intermediate ones) I should their gradients together and replace them.