fix speed 2

run_image_captioning_flax.py

@@ -297,7 +297,7 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
     steps_per_epoch = len(dataset) // batch_size
 
     if shuffle:
-        batch_idx = jax.random.permutation(rng, len(dataset))
+        batch_idx = np.random.permutation(len(dataset))
     else:
         batch_idx = np.arange(len(dataset))
 
@@ -847,7 +847,7 @@ def main():
         if shuffle:
             indices = np.random.permutation(len(train_dataset))
         else:
-            indices = jnp.arange(len(ds))
+            indices = np.arange(len(ds))
 
         for idx in range(num_splits):