Update training

src/run_wav2vec2_pretrain_flax.py

@@ -1,10 +1,3 @@
-import logging
-import sys
-import time
-from dataclasses import field
-from pathlib import Path
-from typing import Dict, List, Optional, Union
-
 # !/usr/bin/env python
 # coding=utf-8
 # Copyright 2021 The HuggingFace Team All rights reserved.
@@ -24,6 +17,12 @@ from typing import Dict, List, Optional, Union
 Training the library models for Wav2Vec.
 """
 # You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
+import logging
+import sys
+import time
+from dataclasses import field
+from pathlib import Path
+from typing import Dict, List, Optional, Union
 
 import numpy as np
 from datasets import DatasetDict, load_dataset
@@ -33,6 +32,7 @@ import flax
 import jax
 import jax.numpy as jnp
 import librosa
+import torchaudio
 import optax
 from flax import jax_utils, traverse_util
 from flax.training import train_state