Iterative solver within loop

If I have a minimzation function which produces some features A,b from some input X for a linear system of equations, I would be able to use torch.solve
to solve the system of equation. However, I have non-linear system of equation which needs to be solved iteratively in the middle. I can optimize this function(with autograd) indiviually using SGD. Is there a way I can combine the iterative loop to the main optimization function?

x = input
A,b = f_1(x)
z = iterative_solve(f_2(A,b))
loss = z' - z