Applying a function along all indices on an axis in PyTorch

I'm trying to implement the Wasserstein Loss function in PyTorch, and I'm referencing the Scipy implementation for this. Since using PyTorch functions within the forward() method implies not having to write the backward() function, I have done this in my code (with the Scipy version just having the equivalent Numpy functions involved). Here is what I have:

class WassersteinLoss(nn.Module): 
  def __init__(self):
    super(WassersteinLoss, self).__init__()
  def forward(self,u,v):
    result = torch.empty((len(u)))
    for i in range(len(u)):
      u_values,v_values = u[i],v[i]
      u_sorter,v_sorter = torch.argsort(u_values),torch.argsort(v_values)
      all_values = torch.cat((u_values,v_values))
      all_values,idx = torch.sort(all_values)

      # Compute the differences between pairs of successive values of u and v.
      deltas = torch.sub(all_values[1:],all_values[:-1])

      # Get the respective positions of the values of u and v among the values of
      # both distributions.
      u_cdf_indices = torch.searchsorted(u_values[u_sorter],all_values[:-1],right=True)
      v_cdf_indices = torch.searchsorted(v_values[v_sorter],all_values[:-1],right=True)

      # Calculate the CDFs of u and v
      u_cdf = torch.div(u_cdf_indices,len(u_values))
      v_cdf = torch.div(v_cdf_indices,len(v_values))

      # Compute the value of the integral based on the CDFs.
      result[i] = torch.sum(torch.multiply(torch.abs(u_cdf-v_cdf),deltas))
    return result.mean()

In the above function, u and v are the vectors of shape (NxM), where N is the number of samples in the batch. Since my for loop essentially goes over all the samples, the calculations on each of them are independent, since the samples in a batch depend on each other. I believe I'd see a significant speedup if I can do away with this for loop. Until now, I have tried performing all computations along the dim=1 axis, but this does not work.

Here is the code for a test case:

from scipy.stats import wasserstein_distance
import torch
import torch.nn as nn

print(wasserstein_distance([0, 1, 3], [5, 6, 8]))
#Output is 5

criterion = WassersteinLoss()
print(criterion(torch.tensor([[0, 1, 3]]), torch.tensor([[5, 6, 8]])))
#Output is tensor(5.)

Any inputs on how to modify the forward() function to eliminate the for loop would be greatly appreciated.