import torch
def normalization(x):
          return (x-x.mean()/x.std())

a=torch.randn(2,2,requires_grad=True)
b=normalization(a)
print(b.mean(),b.std())

If you want forward just create a nn.Module and put this in the forward method. The __init__ can be empty