Support for LayerNorm

I was trying to extend a Vision Transformer model using backpack. However, I encounter the following error:

> UserWarning: Extension saving to grad_batch does not have an extension for Module although the module has parameters

I know that torch.nn.BatchNormNd leads to ill-defined first-order quantities and hence it is not implemented here. Does the same hold for Layer Normalization?

Thank you in advance!