To me, what you are doing is a gradient flow. I would suggest you to have a look at the very good post on Wasserstein gradient flow here : https://www.kernel-operations.io/geomloss/_auto_examples/comparisons/plot_gradient_flows_2D.html#sphx-glr-auto-examples-comparisons-plot-gradient-flows-2d-py 2
In the gradient flow function, they basically do:
for i in range(Nsteps):
        # Compute loss and gradient
        L_αβ = loss(x, y)
        [g]  = torch.autograd.grad(L_αβ, [x])
        x.data -= learning_rate * len(x) * g

This way, you do not need to define an optimizer as you do it yourself ! Hope this helps