move normalization + logit_scale into each model or maybe perhaps some wrapper that does it for all models, think about this more
move normalization + logit_scale into each model or maybe perhaps some wrapper that does it for all models, think about this more