How to customize the loss function using Wasserstein distance in Pytorch?

Question

I am new to using Pytorch. I have two sets of observational data Y and X, probably having different dimensions. My task is to train a function