use jax rng

run_image_captioning_flax.py

@@ -297,7 +297,8 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
     steps_per_epoch = len(dataset) // batch_size
 
     if shuffle:
-        batch_idx = np.random.permutation(len(dataset))
+        batch_idx = jax.random.permutation(rng, len(dataset))
+        batch_idx = np.asarray(batch_idx)
     else:
         batch_idx = np.arange(len(dataset))
 
@@ -845,7 +846,8 @@ def main():
         num_splits = steps // steps_per_split + int(steps % steps_per_split > 0)
 
         if shuffle:
-            indices = np.random.permutation(len(train_dataset))
+            indices = jax.random.permutation(rng, len(train_dataset))
+            indices = np.asarray(indices)
         else:
             indices = np.arange(len(ds))