Looks a bit complicated to me. If you can write your objective function with x, y, z variables then we can use obj_fn(x, y.detach(), z) , it is like giving up code reusability for bug free gradient flow.