Normalize(has epsilon) a tensor along a specific dim

a = torch.Tensor(3,4,5)
for i in range(3):            

I wonder if we can accelerate it a little bit.

a = a * (1 / (a.norm(2, 0, keepdim=True) + 1e-8))

By the way, your original code doesn’t work since it doesn’t save the result of multiplication. You should either assign it back to a[i] or use mul_

Nice catch! Thank you!