Hi @tom
Wow !!!
That’s pretty amazing how quickly you did that !!! Very, very impressive.
Let me go over it and try to do some testing - I need to get my slow old brain working   !! Be nice to do a test verses the unregularized version, i.e.
ot.lp.emd : Unregularized OT
So now that you have a Wasserstein loss that you can backprop through maybe you want to train a plain vanilla GAN with it, it would be interesting to see how it compares, here’s a quite basic version that’s new, and converges quite fast,

  
      github.com
  
  
    AjayTalati/generative-models/blob/master/GAN/boundary_equilibrium_gan/began_pytorch.py 75
import torch
import torch.nn
import torch.nn.functional as nn
import torch.autograd as autograd
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import os
from torch.autograd import Variable
from tensorflow.examples.tutorials.mnist import input_data


mnist = input_data.read_data_sets('../../MNIST_data', one_hot=True)
mb_size = 32
z_dim = 10
X_dim = mnist.train.images.shape[1]
y_dim = mnist.train.labels.shape[1]
h_dim = 128
cnt = 0


  This file has been truncated. show original

  
  
    
    
  
  


Be interesting if you could use your loss layer to improve it?
Perhaps you want to get in touch with Rémi Flamary, http://remi.flamary.com/ 11, I’m sure he’ll be very impressed and be bursting with ideas for possible collaboration 
Best regards,
Ajay