Solved by jdhao in post #9 
              
              
                Now PyTorch have a normalize function, so it is easy to do L2 normalization for features. Suppose x is feature vector of size N*D (N is batch size and D is feature dimension), we can simply use the following 
import torch.nn.functional as F
x = F.normalize(x, p=2, dim=1)