Yea, my code is already written in a way where that will involve a lot of extra bookkeeping.
I found something that works though (simplified from my code):
# Create two distinct optimizers
f_opt = create_optimizer(f.parameters())
g_opt = create_optimizer(g.parameters())

# Forward pass
z = g(x)
y = f(z)

# Update parameters of g only by computing gradient wrt. z first. Gradient on f's params should be None after this.
main_obj = obj_fn(x, y)
z_grad = torch.autograd.grad(main_obj, z, retain_graph = True)
torch.autograd.backward(z, z_grad, retain_graph = True)
g_optimizer.step()

# Update parameters of f only by doing a backwards pass of the entropy objective. Gradient on g's params are not None going into this, but because f_optimizer is only over g's params, that doesn't matter.
h_obj = -entropy(y)
h_obj.backward()
f_optimizer.step()

And I am pretty sure this works (verified which params have None grads after which steps).