Hi @tom,
Many thanks for making this implementation available to the community, I appreciate it a lot. I have almost figured out how to use it, but I have a problem, please forgive me if this is a very naive question.
After instantiating a criterion based on your WassersteinLossStab() class, if I try to compute the loss with a batch size of 1 it seems to work pretty consistently, but if I change the batch size to >1, it crashes:
batch_size = 1 # setting this to 2 breaks the loss
n_samples = 10

x=numpy.arange(n_samples,dtype=numpy.float32)
M=(x[:,numpy.newaxis]-x[numpy.newaxis,:])**2
M/=M.max()

preds = torch.FloatTensor(batch_size,n_samples).uniform_()
preds/=preds.sum(dim=1).unsqueeze(1)
preds = preds.float()

targets = torch.randn(batch_size, n_samples)>0
targets = targets.float()

criterion = WassersteinLossStab(torch.from_numpy(M), lam=0.1)

criterion(preds, targets)

Do you have any clue on what am I doing wrong here? Thanks!