I am trying to implement the paper 'Learning with Wasserstein loss' (the link is https://arxiv.org/abs/1506.05439) then, more specifically, I try to implement the algorithm 1 at page 4 in the paper. However, I don't know how to reflect the result of algorithm 1 to my model via optimizer.step(), e.g. SGD.
For example, when we calculate loss with pytorch, then we can progress the learning with such a following code.
optimizer.zero_grad()
loss.backward()
optimizer.step()
However, the algorithm 1 output the gradient of wasserstein loss with entropic reguralization. Therefore we cant't update it unlike the case of caliculating a loss.
How to update my model by algorithm 1?