Hi @Bauke_Brenninkmeijer,
I created a github repo for this: https://github.com/feevos/pytorch_stuff/blob/master/pytorch_stuff/Notebooks/OMDAdam_WGAN-GP_Gaussians.ipynb 37
the basic idea, from what I recall from the paper, is that you need to perform calculation of the gradients twice. So I tried to implement Algorithm 3, page 20 in the appendix. Where I used the function closure to implement the additional gradient step. So I am calculating the loss twice, sort of.
Perhaps my code has errors, please do check and let me know if you disagree/find anything. But it works well in my tests. Please, if you play with the notebook, let me know what you think of the performance etc. Any comments most appreciated.
I hope you’ll find this useful,
Regards,
Foivos