I tried implementing it this way, but when printing xi it shows that lamb still isn’t optimized for…
It: 0/200000, Loss: 1.04499, xi: [0.][-1.][1.]				 
It: 1/200000, Loss: 0.84020, xi: [0.][-1.][1.]				 
It: 2/200000, Loss: 0.66697, xi: [0.][-1.][1.]				 
It: 3/200000, Loss: 0.52373, xi: [0.][-1.][1.]				 
It: 4/200000, Loss: 0.40832, xi: [0.][-1.][1.]				 
It: 5/200000, Loss: 0.31867, xi: [0.][-1.][1.]				 
It: 6/200000, Loss: 0.25264, xi: [0.][-1.][1.]				 
It: 7/200000, Loss: 0.20781, xi: [0.][-1.][1.]				 
It: 8/200000, Loss: 0.18126, xi: [0.][-1.][1.]				 
It: 9/200000, Loss: 0.16949, xi: [0.][-1.][1.]				 
It: 10/200000, Loss: 0.16853, xi: [0.][-1.][1.]				 
It: 11/200000, Loss: 0.17435, xi: [0.][-1.][1.]				 
It: 12/200000, Loss: 0.18329, xi: [0.][-1.][1.]				 
It: 13/200000, Loss: 0.19237, xi: [0.][-1.][1.]				 
It: 14/200000, Loss: 0.19953, xi: [0.][-1.][1.]
...

I got a hackaround by optimizing for the whole vector xi and manually setting the gradients of the other components to 0 after loss.backward() like so:
xi = nn.Parameter(torch.tensor([[0], [-1], [1]], requires_grad=True, device=device, dtype=torch.float32))
params = [
    {'params': net.parameters(), 'lr': 1e-3},
    {'params': xi, 'lr': 1e-3}
]

optimizer = Adam(params)
while(training):
    # some stuff and computing loss...
    loss.backward()

    xi.grad[1:3] = 0
    optimizer.step()

But this doesn’t really feel right…