I simplified the above code into something more concise that shows what I am trying to do and also shows that it is not happening in pytorch.
By my hand calculation, the second derivative of this at the bottom print statement should be -12.xx but I am getting the first order derivative instead of the second even though I have set create_graph=True. Am I doing something wrong here?
def maml_simulation() -> None:
    x_i = torch.tensor(3.0, requires_grad=True)
    x_j = torch.tensor(4.0, requires_grad=True)
    theta = torch.tensor(2.0, requires_grad=True)
    theta_two = theta.clone()
    loss_func = lambda x: (x - 10) ** 2
    print(f"before training: theta: {theta}")

    prediction = theta_two * x_i
    loss_one = loss_func(prediction)

    grad = torch.autograd.grad(loss_one, theta_two, create_graph=True)[0]
    theta_two -= 0.01 * grad
    print(
        f"first prediction: {prediction}, loss: {loss_one}, grad: {grad}, theta after update: {theta_two}"
    )

    prediction = theta_two * x_j
    loss = loss_func(prediction)
    grad = torch.autograd.grad(loss, theta_two)[0]
    print(f"second prediction: {prediction}, loss: {loss}, second grad: {grad}")