Applying torch.dist across single axis

1

I have two 10000 by 10 arrays:

A = torch.rand((10000, 10))
B = torch.rand((10000, 10))

I want a 10000-dimensional vector with the i_{th} entry being L2-norm of A[i,:] (which is of course a vector of dimension 10) - B[i,:].

What I get with torch.dist(A,B) is a single number summed over both axes.

It does not look from the docs like there is a simple “axis=” keyword to specify which axis I want to sum over.

Any suggestions?

Thanks!