Normalize(has epsilon) a tensor along a specific dim

1 
a = torch.Tensor(3,4,5)
for i in range(3):            
    a[i].mul(1/(torch.norm(a[i],2)+1e-8))

I wonder if we can accelerate it a little bit.

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

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_

4

Nice catch! Thank you!