Thanks for your replies.
Ok so i have something that looks like your init.
x = torch.FloatTensor(self.inputs).to(device)
x.requires_grad = True

This goes straigth into the network which tries to produce the displacement d = x - X.
So far only pytorch’s ops have been used.
Since we get F = dx/dX = grad(d) + I, I naively though i could use d.backward(one)
The problem might comes from the init since this is the only part that might not match perfectly what you just described.