looking for an equivalent of Tensorflow normalization layer in Pytorch
I was using 'tf.keras.layers.experimental.preprocessing.Normalization'. This layer is cool since you can save weights in this layer to normalize any input data to this layer.
However, I couldn't find any normalization layer in Pytorch.
Is there a layer that functions the same role?

There is no built-in that achieves this is PyTorch. However, you can measure the mean and standard deviation yourself (keeping only the relevant axes), then use torchvision.transform.Normalize with those statistics.
For instance in order to measure mean and std over the channels:
&gt;&gt;&gt; x = torch.rand(16, 3, 10, 10)
&gt;&gt;&gt; mean, std = x.mean((0, 2, 3)), x.std((0, 2, 3))
(tensor(0.4941), tensor(0.2899))

Then initialize a transform:
&gt;&gt;&gt; t = torchvision.transform.Normalize(mean, std)

You can use this function on a new dataset to normalize it based on the initial dataset's statistics:
&gt;&gt;&gt; z_normalized = t(z)

