text: from torch import nn
  
  
  class Normalization(nn.Module):
      def __init__(self, mean, std):
          super(Normalization, self).__init__()
          # .view the mean and std to make them [C x 1 x 1] so that they can
          # directly work with image Tensor of shape [B x C x H x W].
          # B is batch size. C is number of channels. H is height and W is width.
          self.mean = mean.clone().detach().view(-1, 1, 1)
          self.std  = std.clone().detach().view(-1, 1, 1)
  
      def forward(self, img):
          # normalize img
          return (img - self.mean) / self.std