In short, I want to enable “create_graph” when doing loss.backward() while using torch.nn layers, such that I can do param.backward() to get the gradient of the final weights w.r.t a hyper parameter.
In details, I am implementing an algorithm to solve a bilevel problem (two nested problems). One can look at the parameters optimised in the inner one as the weights of a torch.nn based model, while the parameters of the outer problem as the hyper parameters. I want to optimise both with gradient descent. Thus, to do one update on the hyper parameters, I need the gradient of the model’s weights (after being trained) with respect to these hyper parameters. This is because the loss function related to the hyper parameters optimisation is a function of the trained model weights.
The problem is that even when I set (create_graph = True) when I backward the inner loss,  optimizer.step() performs in-place updates, so the graph cannot be created. Similarly when replacing optimizer.step() with manually doing the updates on the model weights, as it is still in-place updates:
        for name, param in model.named_parameters():
            param.data =param.data - param.grad

A simplified code of what I want to do:
for t in range(OUT_MAX_ITR):
    model.train()
    for i in range(IN_MAX_ITR):
        optimizer.zero_grad()
        outputs = model(xtr)
        loss = compute_loss(outputs)
        loss.backward(create_graph=True)
        optimizer.step()
 
    theta.grad = None
    out_loss = function_of_model_weights()
    out_loss.backward()
    update_theta(theta, theta.grad)

Here theta is the hyper parameter to be optimised.
Is there a way or a work around in torch to do that second order differentiation (or bilevel optimisation) when working with torch.nn?