I have a model architecture which is a composition of two other models with distinct parameters. So basically my model can be written as y = f(g(x)) where f and g are modules with different sets of parameters.
I have two objectives during learning. One is the main learning objective, and the other is an entropy regularization term. I want to apply these two objectives simultaneously (i.e., on the same inputs with the same state of the network parameters) but I want the main objective to apply only to g's parameters, and the entropy to apply only to f's parameters.
I can set up my code in a way where I can apply the main objective to g only and the entropy to both f and g:
# Optimizer gets parameters from both g and f.
optimizer = create_optimizer(g, f)

# Forward pass and compute the objectives.
z = g(x)
y = f(z)
main_obj = obj_fn(x, y)
h_obj = -entropy(y)

# Compute gradient of main obj wrt. intermediate tensor(s) z, then apply that loss to the paramgers of g.
z_grad = torch.autograd.grad(main_obj, z, retain_graph = True)
torch.autograd.backward(z, z_grad)

# Optimize the entropy wrt. the entire network.
h_obj.backward()

# Update the parameters
optimizer.step()

How can I instead apply the entropy objective only to f's parameters? I know that I can split off z into a version of z detached from g's graph, and then when I optimize the entropy on it, the gradients will only apply to the parameters of f:
z_detach = z.detach()
y_alternative = f(z_detach)
h_obj = -entropy(y_alternative)
h_obj.backward()

but that would require a lot of extra bookkeeping and make the code messy. Is there an easier way I can just compute and apply the gradients over f's parameters?
Thanks!