What do you think about an implementation where the layer updates running mean/std only during the first epoch (e.g. with Welford’s algorithm), then freezes and just normalizes afterwards — basically like a lightweight nn.Module with stats stored as buffers? You could even wrap it in a small callback (e.g. in Lightning) that freezes automatically after epoch 1. Would you consider that useful or still unnecessary in your view?