Simple Example: Solving Lagrange Multiplier with PyTorch

pytorch_lagrange_multi.py

import torch

x = torch.tensor(0, requires_grad=True, dtype=torch.float64)
y = torch.tensor(0, requires_grad=True, dtype=torch.float64)
l = torch.tensor(0, requires_grad=True, dtype=torch.float64)
lr = 0.1

# min x^2+y^2 s.t x+y = 1
for i in range(100):
    L = x**2 + y**2 + l*(1-x-y)
    L.backward()
    x.data.add_(-lr*x.grad.data)
    y.data.add_(-lr*y.grad.data)
    x.grad.detach_()
    x.grad.zero_()
    y.grad.detach_()
    y.grad.zero_()
    l.grad.detach_()
    l.grad.zero_()
    L = x**2 + y**2 + l*(1-x-y)
    L.backward()
    l.data.add_(lr*l.grad.data)
    x.grad.detach_()
    x.grad.zero_()
    y.grad.detach_()
    y.grad.zero_()
    l.grad.detach_()
    l.grad.zero_()

print(x, y, x+y)