[WIP] Flax Speech Recognition Seq2Seq
This PR adds the example script `run_flax_speech_recognition_seq2seq.py`, which can be used to fine-tune any [`Speech Sequence-to-Sequence Model`](https://huggingface.co/docs/transformers/master/en/model_doc/auto#transformers.AutoModelForSpeechSeq2Seq) for automatic speech recognition **in Flax** on one of the [official speech recognition datasets](https://huggingface.co/datasets?task_ids=task_ids:automatic-speech-recognition) or a custom dataset.

What remains TODO:
- [x] Add an adapter layer when loading a pre-trained Wav2Vec2 encoder checkpoint through a keyword argument. See: https://github.com/huggingface/transformers/pull/16056
- [x] Flax length grouped sampler: group training data into mega batches sorted by input length in order to minimise the amount of padding required. Place the batch of `batch_size` containing the element of maximum length placed first, so that an OOM happens sooner rather than later. See: https://github.com/huggingface/transformers/pull/16058/commits/d66c59de31ce3d8c61127b63b9198b69a0930ccd
- [x] Fix compilation issues with the Flax DataCollator: once the training data is grouped by input length, pad the input sequences using `pad_to_multiple_of`. Independently pad the target labels using `max_length`. See: https://github.com/huggingface/transformers/pull/16058/commits/d5a5fbdfeca67275e75c55498af0a3282c1c0741
- [ ] Verify the script works and that it yields equivalent results to PyTorch. Fine-tune Wav2Vec2-2-Bart-large-cnn on `librispeech-asr` in order to reproduce the PyTorch training results at https://huggingface.co/sanchit-gandhi/wav2vec2-2-bart-large-cnn/tensorboard.
- [ ] Save the feature encoder outputs during/before the first epoch and use these as the model inputs for subsequent epochs.
- [x] Gradient accumulation using the Optax optimiser library functions (`apply_every` or `MultiSteps`). Huge memory overhead. Switch to a custom implementations (see below).
- [x] Custom gradient accumulation. See: https://github.com/huggingface/transformers/pull/16058/commits/3ef71d141fb9b643959326c123026c454c4efa13
- [x] Gradient Checkpointing (aborted) - requires an `nn.remat` module in every Flax model layer, and shown not to give huge improvements in the Dalle-mini project.
- [ ] (Load and Continue from Last Checkpoint)
Create a tiny model:
```python
import jax.numpy as jnp
from transformers import AutoFeatureExtractor, AutoTokenizer, FlaxSpeechEncoderDecoderModel


encoder_id = "hf-internal-testing/tiny-random-wav2vec2"
decoder_id = "hf-internal-testing/tiny-random-bart"

model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained(
    encoder_id, decoder_id, encoder_from_pt=True, decoder_from_pt=True
)

model.config.encoder.feat_proj_dropout = 0.0
model.config.encoder.final_dropout = 0.0
model.config.encoder.mask_time_prob = 0.1
model.config.decoder_start_token_id = model.config.decoder.bos_token_id
model.config.pad_token_id = model.config.decoder.pad_token_id
model.config.eos_token_id = model.config.decoder.eos_token_id
model.config.max_length = 20
model.config.num_beams = 1
model.config.encoder.layerdrop = 0.0
model.config.use_cache = False
model.config.processor_class = "Wav2Vec2Processor"

# check if generation works
out = model.generate(jnp.ones((1, 2000)))

model.save_pretrained("./")

feature_extractor = AutoFeatureExtractor.from_pretrained(encoder_id)
feature_extractor.save_pretrained("./")
tokenizer = AutoTokenizer.from_pretrained(decoder_id)
tokenizer.save_pretrained("./")
```
Train on dummy dataset:
```bash
#!/usr/bin/env bash
python run_flax_speech_recognition_seq2seq.py \
        --dataset_name="hf-internal-testing/librispeech_asr_dummy" \
        --model_name_or_path="./" \
        --output_dir="./" \
        --dataset_config_name="clean" \
        --train_split_name="validation" \
        --eval_split_name="validation" \
        --preprocessing_num_workers="1" \
        --length_column_name="input_length" \
        --overwrite_output_dir \
        --num_train_epochs="2" \
        --per_device_train_batch_size="2" \
        --per_device_eval_batch_size="2" \
        --gradient_accumulation_steps="2" \
        --generation_max_length="40" \
        --generation_num_beams="1" \
        --learning_rate="3e-4" \
        --warmup_steps="5" \
        --evaluation_strategy="steps" \
        --text_column_name="text" \
        --freeze_feature_encoder \
        --predict_with_generate \
        --do_lower_case \
        --do_train \
        --do_eval
```<|||||>_The documentation is not available anymore as the PR was closed or merged._<|||||>Let's continue here without gradient_checkpointing for now. Don't think that this should block us for the first Wav2Vec2-Bart training run<|||||>Does it work with `gradient_checkpointing`? Will do a proper review on Monday, but in the mean time feel free to already go ahead and try it out with Wav2Vec2-Bart on 8 TPUs :-) <|||||><details>
<summary>Create 'tiny' Wav2Vec2-2-Bart model (`tied_word_embeddings=False`)</summary>

```python
import numpy as np
from transformers import AutoFeatureExtractor, AutoTokenizer, FlaxSpeechEncoderDecoderModel


encoder_id = "hf-internal-testing/tiny-random-wav2vec2"
decoder_id = "hf-internal-testing/tiny-random-bart"

model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained(encoder_id, decoder_id, encoder_from_pt=True, decoder_from_pt=True, 
encoder_add_adapter=True, decoder_tie_word_embeddings=False)

model.config.decoder_start_token_id = model.config.decoder.bos_token_id
model.config.pad_token_id = model.config.decoder.pad_token_id
model.config.eos_token_id = model.config.decoder.eos_token_id
model.config.max_length = 20
model.config.num_beams = 1
model.config.use_cache = False
model.config.processor_class = "Wav2Vec2Processor"

# check if generation works
out = model.generate(np.ones((1, 2000)))

model.save_pretrained("./")

feature_extractor = AutoFeatureExtractor.from_pretrained(encoder_id)
feature_extractor.save_pretrained("./")
tokenizer = AutoTokenizer.from_pretrained(decoder_id)
tokenizer.save_pretrained("./")
```
</details>

<details>
<summary>Run on 5% of Librispeech ASR on TPU with float32 (full) precision</summary>

```
#!/usr/bin/env bash
JAX_DEFAULT_MATMUL_PRECISION=float32 python run_flax_speech_recognition_seq2seq.py \
        --dataset_name="librispeech_asr" \
        --model_name_or_path="./" \
        --dataset_config_name="clean" \
        --train_split_name="train.100[:5%]" \
        --eval_split_name="validation[:5%]" \
        --output_dir="./" \
        --preprocessing_num_workers="16" \
        --length_column_name="input_length" \
        --overwrite_output_dir \
        --num_train_epochs="1" \
        --per_device_train_batch_size="2" \
        --per_device_eval_batch_size="2" \
        --logging_steps="1" \
        --max_duration_in_seconds="10" \
        --max_target_length="32" \
        --generation_max_length="40" \
        --generation_num_beams="1" \
        --learning_rate="3e-4" \
        --warmup_steps="500" \
        --text_column_name="text" \
        --save_total_limit="1" \
        --freeze_feature_encoder \
        --predict_with_generate \
        --do_lower_case \
        --do_eval \
        --do_train 

```

</details>


<details>
<summary>'Tiny' model logs: l2-norm of encoder and decoder hidden-states (on a per-layer basis). Real-valued outputs on the first forward pass, backward pass and proceeding training steps.</summary>

```
03/18/2022 14:02:38 - INFO - __main__ - ***** Running training *****
03/18/2022 14:02:38 - INFO - __main__ -   Num examples = 94
03/18/2022 14:02:38 - INFO - __main__ -   Num Epochs = 1
03/18/2022 14:02:38 - INFO - __main__ -   Instantaneous batch size per device = 2
03/18/2022 14:02:38 - INFO - __main__ -   Total train batch size (w. parallel & distributed) = 16
03/18/2022 14:02:38 - INFO - __main__ -   Total optimization steps = 5
Step... (1) | decoder_hidden_states: [31.988552 31.934507 31.932938] | encoder_hidden_states: [121.05971 121.14229 121.0128  121.08176 121.19736] | encoder_outputs: 0.053542863577604294 | inputs: 337.59503173828125 | logits: 19.7552
03247070312 | loss: 6.910846710205078 |                                                                                                                                                                                                 
Step... (2) | decoder_hidden_states: [31.978403 31.85397  31.924433] | encoder_hidden_states: [120.8189   120.81773  120.836624 120.75368  120.73546 ] | encoder_outputs: 0.053248822689056396 | inputs: 332.2764587402344 | logits: 19.
764877319335938 | loss: 6.913311958312988 |                                                                                                                                                                                             
Step... (3) | decoder_hidden_states: [32.03605  32.0115   31.969864] | encoder_hidden_states: [120.74285  120.60301  120.512054 120.62747  120.67913 ] | encoder_outputs: 0.05340719223022461 | inputs: 345.63580322265625 | logits: 19.
75674057006836 | loss: 6.906001091003418 |                                                                                                                                                                                              
Step... (4) | decoder_hidden_states: [31.977776 31.981033 31.923492] | encoder_hidden_states: [121.15105 120.91719 121.1255  121.18452 121.01183] | encoder_outputs: 0.0538754016160965 | inputs: 369.9171447753906 | logits: 19.7641830
44433594 | loss: 6.903560638427734 |
Training...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:24<00:00,  4.94s/it]
Epoch ... (1/1): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:24<00:00, 24.71s/it]
```
</details><|||||><details>
<summary>Create ‘large’ Wav2Vec2-2-Bart model (`tied_word_embeddings=False`)</summary>

```python
import numpy as np
from transformers import AutoFeatureExtractor, AutoTokenizer, FlaxSpeechEncoderDecoderModel

encoder_id = "hf-internal-testing/tiny-random-wav2vec2"
decoder_id = "hf-internal-testing/tiny-random-bart"

model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained(encoder_id, decoder_id, encoder_from_pt=True, decoder_from_pt=True, 
encoder_add_adapter=True, decoder_tie_word_embeddings=False)

model.config.decoder_start_token_id = model.config.decoder.bos_token_id
model.config.pad_token_id = model.config.decoder.pad_token_id
model.config.eos_token_id = model.config.decoder.eos_token_id
model.config.max_length = 20
model.config.num_beams = 1
model.config.use_cache = False
model.config.processor_class = "Wav2Vec2Processor"

# check if generation works
out = model.generate(np.ones((1, 2000)))

model.save_pretrained("./")

feature_extractor = AutoFeatureExtractor.from_pretrained(encoder_id)
feature_extractor.save_pretrained("./")
tokenizer = AutoTokenizer.from_pretrained(decoder_id)
tokenizer.save_pretrained("./")
```
</details>

<details>
<summary>Run on 5% of Librispeech ASR on TPU with float32 (full) precision - identical bash script to that for 'tiny'</summary>

```
#!/usr/bin/env bash
JAX_DEFAULT_MATMUL_PRECISION=float32 python run_flax_speech_recognition_seq2seq.py \
        --dataset_name="librispeech_asr" \
        --model_name_or_path="./" \
        --dataset_config_name="clean" \
        --train_split_name="train.100[:5%]" \
        --eval_split_name="validation[:5%]" \
        --output_dir="./" \
        --preprocessing_num_workers="16" \
        --length_column_name="input_length" \
        --overwrite_output_dir \
        --num_train_epochs="1" \
        --per_device_train_batch_size="2" \
        --per_device_eval_batch_size="2" \
        --logging_steps="1" \
        --max_duration_in_seconds="10" \
        --max_target_length="32" \
        --generation_max_length="40" \
        --generation_num_beams="1" \
        --learning_rate="3e-4" \
        --warmup_steps="500" \
        --text_column_name="text" \
        --save_total_limit="1" \
        --freeze_feature_encoder \
        --predict_with_generate \
        --do_lower_case \
        --do_eval \
        --do_train 

```

</details>

<details>
<summary>'Large' model logs: l2-norm of encoder and decoder hidden-states (on a per-layer basis). On the first forward pass, the l2-norm of the encoder hidden-states is well defined. The logs show that the first decoder-hidden state returns 'nan' values.</summary>

```
03/18/2022 14:16:05 - INFO - __main__ - ***** Running training *****                                                                                                                                                                    
03/18/2022 14:16:05 - INFO - __main__ -   Num examples = 221                                                                                                                                                                            
03/18/2022 14:16:05 - INFO - __main__ -   Num Epochs = 1                                                                                                                                                                                
03/18/2022 14:16:05 - INFO - __main__ -   Instantaneous batch size per device = 2                                                                                                                                                       
03/18/2022 14:16:05 - INFO - __main__ -   Total train batch size (w. parallel & distributed) = 16                                                                                                                                       
03/18/2022 14:16:05 - INFO - __main__ -   Total optimization steps = 13                                                                                                                                                                 
Step... (1) | decoder_hidden_states: [nan nan nan nan nan nan nan nan nan nan nan nan nan] | encoder_hidden_states: [410571.8  410428.7  410573.44 410535.3  410262.44 409993.22 410005.3                                               
 410596.28 410562.3  410165.12 410397.44 409875.1  410200.16 410837.16                                                                                                                                                                  
 410408.94 410557.22 410525.5  410319.47 410763.8  410527.88 410446.94                                                                                                                                                                  
 410469.2  410431.88 410328.4  410323.3 ] | encoder_outputs: 15.634440422058105 | inputs: 439.39068603515625 | logits: nan | loss: nan |                                                                                                
Step... (2) | decoder_hidden_states: [nan nan nan nan nan nan nan nan nan nan nan nan nan] | encoder_hidden_states: [nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan                                            
 nan nan nan nan nan nan nan] | encoder_outputs: nan | inputs: 411.89910888671875 | logits: nan | loss: nan |    
```<|||||><details>
<summary>Create 'large' Wav2Vec2-2-GPT2 model</summary>

```python
import numpy as np
from transformers import AutoFeatureExtractor, FlaxSpeechEncoderDecoderModel, GPT2Tokenizer

encoder_id = "facebook/wav2vec2-large-lv60"
decoder_id = "gpt2-medium"

model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained(encoder_id, decoder_id, encoder_add_adapter=True)

# force GPT2 to append EOS to begin and end of seq
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
    outputs = [self.bos_token_id] + token_ids_0 + [self.eos_token_id]
    return outputs

GPT2Tokenizer.build_inputs_with_special_tokens = build_inputs_with_special_tokens
gpt2_tokenizer = GPT2Tokenizer.from_pretrained(decoder_id)
# set pad_token_id to unk_token_id, note: unk_token_id == eos_token_id == bos_token_id
gpt2_tokenizer.pad_token = gpt2_tokenizer.unk_token
gpt2_tokenizer.save_pretrained("./")

model.config.pad_token_id = gpt2_tokenizer.pad_token_id
model.config.decoder_start_token_id = model.config.decoder.bos_token_id
model.config.eos_token_id = model.config.decoder.eos_token_id
model.config.max_length = 40
model.config.num_beams = 1
model.config.use_cache = False
model.config.decoder.use_cache = False
model.config.processor_class = "Wav2Vec2Processor"

# check if generation works
out = model.generate(np.ones((1, 2000)))

model.save_pretrained("./")

feature_extractor = AutoFeatureExtractor.from_pretrained(encoder_id)
feature_extractor.save_pretrained("./")
```
</details>

<details>
<summary>Run on 5% of Librispeech ASR on TPU with float32 (full) precision</summary>

```
#!/usr/bin/env bash
JAX_DEFAULT_MATMUL_PRECISION=float32 python run_flax_speech_recognition_seq2seq.py \
        --dataset_name="librispeech_asr" \
        --model_name_or_path="./" \
        --dataset_config_name="clean" \
        --train_split_name="train.100[:5%]" \
        --eval_split_name="validation[:5%]" \
        --output_dir="./" \
        --preprocessing_num_workers="16" \
        --length_column_name="input_length" \
        --overwrite_output_dir \
        --num_train_epochs="1" \
        --per_device_train_batch_size="2" \
        --per_device_eval_batch_size="2" \
        --logging_steps="1" \
        --max_duration_in_seconds="10" \
        --max_target_length="32" \
        --generation_max_length="40" \
        --generation_num_beams="1" \
        --learning_rate="3e-4" \
        --warmup_steps="500" \
        --text_column_name="text" \
        --save_total_limit="1" \
        --freeze_feature_encoder \
        --predict_with_generate \
        --do_lower_case \
        --do_eval \
        --do_train 

```

</details>

<details>
<summary>Wav2Vec2-2-gpt2 'large' model logs: l2-norm of encoder and decoder hidden-states (on a per-layer basis). Real-valued outputs on the first forward pass, backward pass and proceeding training steps.</summary>

```
03/18/2022 14:29:28 - INFO - __main__ - ***** Running training *****                                                                                                                                                                    
03/18/2022 14:29:28 - INFO - __main__ -   Num examples = 231                                                                                                                                                                            
03/18/2022 14:29:28 - INFO - __main__ -   Num Epochs = 1                                                                                                                                                                                
03/18/2022 14:29:28 - INFO - __main__ -   Instantaneous batch size per device = 2                                                                                                                                                       
03/18/2022 14:29:28 - INFO - __main__ -   Total train batch size (w. parallel & distributed) = 16                                                                                                                                       
03/18/2022 14:29:28 - INFO - __main__ -   Total optimization steps = 14                                                                                                                                                                 
Step... (1) | decoder_hidden_states: [  27.705173  984.6688   1069.8226   1278.6012   5160.4707   5699.3545                                                                                                                             
 5908.879    6050.114    6143.709    6212.514    6243.5605   6270.414                                                                                                                                                                   
 6293.3115   6315.0156   6347.5215   6391.46     6452.412    6532.633                                                                                                                                                                   
 6639.2114   6795.3633   6964.0645   7185.9297   7702.9453   7484.8037                                                                                                                                                                  
 1918.5354  ] | encoder_hidden_states: [400464.72 400572.25 400476.56 400581.   400228.7  399510.75 400316.12                                                                                                                           
 400524.7  400276.5  400623.3  400237.38 399885.06 400458.94 400540.94                                                                                                                                                                  
 400962.78 401065.44 400384.72 400907.28 401037.88 400596.56 400616.4                                                                                                                                                                   
 400706.3  400376.38 400392.06 400209.72] | encoder_outputs: 15.65604305267334 | inputs: 441.4826354980469 | logits: 153219.21875 | loss: 5.277106285095215 |                                                                           
Step... (2) | decoder_hidden_states: [  27.50412  982.7307  1082.1393  1266.6243  5323.2876  5468.919                                                                                                                                   
 5817.261   5984.31    6098.169   6172.6177  6210.377   6239.094                                                                                                                                                                        
 6264.5156  6287.0127  6320.5957  6366.2725  6433.3057  6517.029                                                                                                                                                                        
 6629.832   6783.8105  6960.3086  7182.701   7713.1733  7574.8164                                                                                                                                                                       
 1835.7288 ] | encoder_hidden_states: [399864.1  399652.38 399964.25 400088.06 399837.88 399556.25 399368.34                                                                                                                            
 400199.38 399967.75 399810.5  400387.34 399937.6  399366.7  400193.97                                                                                                                                                                  
 400485.3  400250.5  400265.6  400288.16 399987.75 400048.06 399985.75
Step... (3) | decoder_hidden_states: [  27.585026  967.4181   1059.0647   1267.0371   5161.367    5522.994                                                                                                                              
 5914.909    6066.3604   6165.132    6233.869    6263.9297   6289.3525                                                                                                                                                                  
 6310.677    6330.4287   6361.243    6401.865    6460.8145   6535.998                                                                                                                                                                   
 6634.0107   6781.088    6954.6953   7167.119    7668.0586   7605.51                                                                                                                                                                    
 1945.8069  ] | encoder_hidden_states: [395648.75 395167.88 395541.88 395567.1  395223.38 395166.2  395408.7                                                                                                                            
 395719.5  395231.8  395440.8  396132.2  395191.97 395275.75 395832.78                                                                                                                                                                  
 395969.94 396130.06 395876.16 395957.06 395777.8  395564.5  395641.7                                                                                                                                                                   
 395728.62 395525.12 395319.56 395593.94] | encoder_outputs: 15.668608665466309 | inputs: 469.32513427734375 | logits: 153649.0 | loss: 5.771542072296143 |                                                                             
Step... (4) | decoder_hidden_states: [  26.962208  973.65936  1086.2867   1296.3983   5358.009    5720.5425                                                                                                                             
 5965.629    6129.83     6257.0303   6342.369    6381.131    6410.9604                                                                                                                                                                  
 6438.2354   6461.836    6497.192    6543.486    6610.591    6694.203                                                                                                                                                                   
 6808.67     6969.8984   7143.211    7362.4785   7883.0137   7938.967                                                                                                                                                                   
 1861.9117  ] | encoder_hidden_states: [396724.56 396978.3  397126.6  397024.62 396497.25 396432.4  396626.25                                                                                                                           
 396505.7  396774.38 396963.28 396561.06 396375.7  396653.56 396788.94                                                                                                                                                                  
 397315.38 397261.66 397600.88 397143.75 397139.22 397015.8  397229.72                                                                                                                                                                  
 397317.12 397280.7  397117.   396934.78] | encoder_outputs: 15.405157089233398 | inputs: 413.68218994140625 | logits: 150058.84375 | loss: 5.734399795532227 | 
Step... (5) | decoder_hidden_states: [  27.54343  979.775   1076.0803  1301.9713  5163.4297  5679.1533                                                                                                                         [34/1854]
 5905.4297  6039.828   6160.733   6236.451   6273.1816  6301.7246                                                                                                                                                                       
 6326.33    6349.4707  6383.63    6429.527   6494.6816  6576.9424                                                                                                                                                                       
 6689.077   6847.369   7022.1494  7239.0347  7760.078   7641.672                                                                                                                                                                        
 1898.6593 ] | encoder_hidden_states: [406666.4  406645.62 406751.12 406984.3  406793.06 406148.7  406196.3                                                                                                                             
 406840.38 406524.12 406709.25 406924.88 406276.   406272.8  406852.44                                                                                                                                                                  
 406731.38 406783.12 406927.25 406774.06 406759.97 406678.8  406413.12                                                                                                                                                                  
 406399.38 406823.28 406723.03 406760.56] | encoder_outputs: 15.645303726196289 | inputs: 433.8205261230469 | logits: 150390.6875 | loss: 6.245608329772949 |                                                                           
Step... (6) | decoder_hidden_states: [  27.222858  990.08167  1092.7056   1300.3113   5371.709    5725.07                                                                                                                               
 5957.906    6129.3535   6248.9053   6329.303    6366.3906   6395.6523                                                                                                                                                                  
 6422.501    6445.4014   6480.754    6528.621    6596.4595   6680.3633                                                                                                                                                                  
 6796.694    6957.9355   7141.7573   7361.3135   7884.8945   7559.6064                                                                                                                                                                  
 1813.7799  ] | encoder_hidden_states: [400203.62 400314.06 400102.16 400408.5  400243.12 399775.75 399778.84                                                                                                                           
 400241.03 400312.3  400378.88 400467.84 400262.25 400152.66 400762.97                                                                                                                                                                  
 400715.38 400965.47 400972.2  400946.34 400621.56 400494.2  400381.94                                                                                                                                                                  
 400251.25 400610.   400565.16 399859.5 ] | encoder_outputs: 15.666407585144043 | inputs: 389.3349914550781 | logits: 149033.015625 | loss: 5.595775127410889 |                                                                         
Step... (7) | decoder_hidden_states: [  27.487904  977.2353   1077.8563   1276.3774   5150.678    5677.568                                                                                                                              
 5895.86     6052.6934   6171.8306   6248.2627   6286.8447   6314.7285                                                                                                                                                                  
 6340.3857   6362.9355   6397.525    6443.6978   6508.964    6591.554                                                                                                                                                                   
 6704.009    6859.457    7038.1436   7254.1885   7778.0864   7743.4062                                                                                                                                                                  
 1841.6183  ] | encoder_hidden_states: [400044.25 399508.06 399957.38 400000.25 399269.62 399271.3  399618.38                                                                                                                           
 399673.8  399372.75 399543.03 399945.75 399298.53 399452.9  399628.9                                                                                                                                                                   
 399918.3  399923.44 399984.3  400002.2  400012.94 399586.66 399763.53                                                                                                                                                                  
 399891.75 399923.03 399489.12 399915.62] | encoder_outputs: 15.693841934204102 | inputs: 439.5829162597656 | logits: 149439.1875 | loss: 5.356408596038818 |                                                                           
Step... (8) | decoder_hidden_states: [  27.260107  981.9591   1082.9978   1303.3995   5347.8193   5725.1963                                                                                                                             
 5988.419    6157.9927   6286.791    6367.6123   6408.1953   6437.547                                                                                                                                                                   
 6463.213    6486.0117   6520.4287   6567.293    6633.672    6717.458                                                                                                                                                                   
 6832.828    6996.592    7176.0205   7395.158    7912.5283   7769.053                                                                                                                                                                   
 1882.0337  ] | encoder_hidden_states: [381122.   380513.5  381451.97 381533.62 380671.   380749.44 380397.56                                                                                                                           
 381154.94 381103.56 381389.25 381325.7  380440.2  380236.6  380851.62                                                                                                                                                                  
 381388.38 381022.12 381248.9  381925.8  381047.94 380761.06 381296.44                                                                                                                                                                  
 380768.56 381124.38 380821.7  381050.06] | encoder_outputs: 15.531478881835938 | inputs: 428.942626953125 | logits: 151706.375 | loss: 5.479863166809082 |                                                                             
Step... (9) | decoder_hidden_states: [  27.203522  970.3016   1070.8975   1267.3228   5149.606    5495.448                                                                                                                              
 5720.803    5930.7812   6054.88     6134.3813   6173.323    6201.705                                                                                                                                                                   
 6227.37     6249.925    6284.9014   6332.4463   6399.3584   6485.218                                                                                                                                                                   
 6600.2188   6764.827    6945.5586   7167.542    7712.4214   7688.0728
 1921.8232  ] | encoder_hidden_states: [397437.84 397162.2  397749.75 397334.8  397161.62 396930.25 396887.62
 397239.2  397335.62 397099.62 397290.4  397293.7  397000.25 397579.4
 397912.12 397778.8  397601.94 397803.25 397521.16 397613.8  397667.56
 397828.3  397397.75 397247.25 397604.25] | encoder_outputs: 15.765214920043945 | inputs: 418.4049987792969 | logits: 153519.6875 | loss: 5.3556318283081055 |
Step... (10) | decoder_hidden_states: [  27.101421  971.5001   1079.2378   1281.2228   5339.7983   5706.5225        
 5975.9727   6123.8633   6253.355    6338.587    6379.96     6409.5786                                              
 6435.078    6458.2617   6492.785    6539.077    6604.3154   6686.5317
 6799.6772   6959.0547   7134.6025   7348.9043   7864.45     7982.283
 1889.574   ] | encoder_hidden_states: [402673.25 402494.44 402451.8  402582.56 402037.75 401663.47 402210.
 402500.9  402219.3  402320.1  402611.62 402244.44 402333.94 402710.44
 402725.62 402615.78 402902.   403058.94 402817.12 402506.88 402510.16
 402621.66 402858.06 402372.06 402555.2 ] | encoder_outputs: 15.740177154541016 | inputs: 427.6776123046875 | logits: 151765.953125 | loss: 5.4407196044921875 |
Step... (11) | decoder_hidden_states: [  27.101858  960.69806  1057.977    1275.9325   4993.126    5745.655         
 6005.66     6165.303    6280.248    6355.59     6392.543    6419.9526                                              
 6444.1123   6465.8066   6498.427    6542.979    6606.0117   6687.626
 6797.332    6948.961    7114.5967   7316.5913   7814.159    7582.8237
 1949.6685  ] | encoder_hidden_states: [392330.47 392076.62 392226.   392236.4  391949.8  391717.7  392134.75
 392420.97 391965.97 392092.62 392427.   391836.06 392229.9  392766.7
 392375.94 392607.5  392755.4  392962.8  392651.25 392396.16 392297.78
 392674.38 392437.44 392127.75 392342.3 ] | encoder_outputs: 15.853521347045898 | inputs: 443.4962463378906 | logits: 154125.15625 | loss: 4.923087120056152 |
Step... (12) | decoder_hidden_states: [  27.225206  968.9116   1067.0066   1284.6025   5543.5396   5725.6973        
 5977.792    6148.306    6284.0425   6358.6274   6394.204    6420.883                                               
 6445.0244   6466.761    6500.2197   6545.107    6607.967    6688.363
 6796.278    6947.9785   7114.134    7320.1875   7820.391    7790.4746
 1938.8198  ] | encoder_hidden_states: [389132.75 389117.12 389249.44 389016.8  389135.5  388934.06 388708.4
 389077.2  389071.62 389255.16 389186.56 389063.25 388612.8  388886.75
 389398.06 389742.66 389516.25 389593.62 389379.5  389351.7  389328.75
 389318.   389152.44 389027.44 389372.9 ] | encoder_outputs: 15.960776329040527 | inputs: 445.12548828125 | logits: 153455.5625 | loss: 5.53969669342041 |
Step... (13) | decoder_hidden_states: [  27.60433  964.4414  1066.7817  1286.687   5361.041   5721.004              
 5953.537   6131.306   6260.2393  6339.0967  6373.6455  6400.8584                                                   
 6425.586   6447.6943  6481.8457  6526.724   6591.548   6674.255
 6781.3926  6929.332   7095.1914  7306.254   7807.132   7699.7607
 1906.3538 ] | encoder_hidden_states: [391200.66 390364.25 390781.3  390942.88 390509.3  390213.06 390245.28
 390339.44 390532.6  390483.56 390849.3  390550.12 390550.4  391182.88
 390956.2  391014.22 390799.22 391230.6  391087.06 390722.16 390951.2
 391262.56 391213.2  390638.3  390842.12] | encoder_outputs: 16.03491973876953 | inputs: 460.640380859375 | logits: 150500.109375 | loss: 5.521194934844971 |
Training...: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 14/14 [04:26<00:00, 19.02s/it]
Epoch ... (1/1): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [04:26<00:00, 266.27s/it]
```
</details><|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.<|||||>Training is moved to https://github.com/sanchit-gandhi/seq2seq-speech . Should we close this one for now @sanchit-gandhi ?