deepspeech improvements / bug fixes :
deepspeech improvements / bug fixes :

1) We were missing activation function in our FeedForwardModule layer

2) There is an additional BatchNorm in convolution subsampling layer and it's unclear if it should be present, conformer doesn't use it and in original deepspeech paper it's unclear if it's present ... this extra BatchNorm causes slow JIT compilation due to XLA trying to constant fold this op.

3) BatchNorm layer currently doesn't use cross device synced mean and variance so enabling that option to see if less noisier stats help training.

PiperOrigin-RevId: 478878177

google/init2winit