🚀 The feature, motivation and pitch

For some models the best normalization method is using weight_norm. This is especially true when training for a hardware implementation for processing high resolution images. It's not possible to use batch_norm because the batch size is one, group/layer normalization adds calculations and latency to the hardware implementation.

Alternatives

Allowing fallback on the CPU during training with the obvious training time penalty

Additional context

No response

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev