基于pytorch实现牛顿迭代法求函数极小值

import torch
from torch.autograd import Variable

def f(x):
    y = x ** 2
    return y

x = Variable(torch.Tensor([5]), requires_grad=True)

for i in range(100):
    grad_x = torch.autograd.grad(f(x), x, create_graph=True)
    grad_grad_x = torch.autograd.grad(grad_x[0], x)
    x = Variable(x.data - grad_x[0].data / grad_grad_x[0].data, requires_grad=True)
print(x)