I think the most important one to implement first is BatchNorm1d which does batch normalization over a 2D or 3D input.
We can then move to higher dimensions like BatchNorm2d/3d when needed.
I have a rudimentary BatchNorm layer which can only do eval in my resnet implementation which should be factored out and made into a full fledged layer
Pytorch has a few normalization layers that ChAI does not currently implement. See https://pytorch.org/docs/stable/nn.html#normalization-layers for a list of some of the layers.
I think the most important one to implement first is BatchNorm1d which does batch normalization over a 2D or 3D input. We can then move to higher dimensions like BatchNorm2d/3d when needed.