Applying l2 normalization and recieve the same shape as input tensor

I have already seen some topic about the normalization and no one include my problem.

I have a code in Keras which applies l2 normalization on a matrix and returns the result with the same shape of input:
K.l2_normalize(input, axis=0)

However, It seems that torch.norm function reduces the dimension of input tensor. Is there any equivalent keras norm function in the pytorch or should I implement it from scratch?
I want to have a l2 normalized tensor with the same shape as input

Thank you in advance.

Pytorch has torch.nn.MSELoss() and the functional torch.nn.functional.mse_loss().

Passing reduce=False to any of these will keep your dimensions:

a = torch.rand([2, 3, 20, 20])
b = torch.rand([2, 3, 20, 20])

torch.nn.functional.mse_loss(a, b, reduce=False).shape

> torch.Size([2, 3, 20, 20])