fix speed

run_image_captioning_flax.py

@@ -299,7 +299,7 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
     if shuffle:
         batch_idx = jax.random.permutation(rng, len(dataset))
     else:
-        batch_idx = jnp.arange(len(dataset))
+        batch_idx = np.arange(len(dataset))
 
     batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.
     batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
@@ -845,7 +845,7 @@ def main():
         num_splits = steps // steps_per_split + int(steps % steps_per_split > 0)
 
         if shuffle:
-            indices = jax.random.permutation(input_rng, len(ds))
+            indices = np.random.permutation(len(train_dataset))
         else:
             indices = jnp.arange(len(ds))
 
@@ -864,7 +864,7 @@ def main():
                 "test": "prediction",
             }
 
-            _ds =_ds.map(
+            _ds = _ds.map(
                 feature_extraction_fn,
                 batched=True,
                 num_proc=data_args.preprocessing_num_workers,