Use the GroupNorm 126 as followed:
nn.GroupNorm(1, out_channels)
It is equivalent with LayerNorm. It is useful if you only now the number of channels of your input and you want to define your layers as such
nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size, stride), nn.GroupNorm(1, out_channels), nn.ReLU())