fix

run_image_captioning_flax.py

@@ -849,12 +849,12 @@ def main():
         if "train" not in dataset:
             raise ValueError("--do_train requires a train dataset")
         train_dataset = dataset["train"]
-        # remove problematic examples 
+        if data_args.max_train_samples is not None:
+            train_dataset = train_dataset.select(range(data_args.max_train_samples))
+        # remove problematic examples
         # (if feature extraction is performed at the beginning, the filtering is done during preprocessing not here)
         if not run_feat_ext_at_beginning:
             train_dataset = train_dataset.filter(filter_fn, batched=True, num_proc=data_args.preprocessing_num_workers)
-        if data_args.max_train_samples is not None:
-            train_dataset = train_dataset.select(range(data_args.max_train_samples))
         train_dataset = train_dataset.map(
             function=function_kwarg,
             batched=True,
@@ -874,12 +874,12 @@ def main():
         if "validation" not in dataset:
             raise ValueError("--do_eval requires a validation dataset")
         eval_dataset = dataset["validation"]
-        # remove problematic examples 
+        if data_args.max_eval_samples is not None:
+            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
+        # remove problematic examples
         # (if feature extraction is performed at the beginning, the filtering is done during preprocessing not here)
         if not run_feat_ext_at_beginning:
             eval_dataset = eval_dataset.filter(filter_fn, batched=True, num_proc=data_args.preprocessing_num_workers)
-        if data_args.max_eval_samples is not None:
-            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
         eval_dataset = eval_dataset.map(
             function=function_kwarg,
             batched=True,
@@ -899,12 +899,12 @@ def main():
         if "test" not in dataset:
             raise ValueError("--do_predict requires a test dataset")
         predict_dataset = dataset["test"]
-        # remove problematic examples 
+        if data_args.max_predict_samples is not None:
+            predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
+        # remove problematic examples
         # (if feature extraction is performed at the beginning, the filtering is done during preprocessing not here)
         if not run_feat_ext_at_beginning:
             predict_dataset = predict_dataset.filter(filter_fn, batched=True, num_proc=data_args.preprocessing_num_workers)
-        if data_args.max_predict_samples is not None:
-            predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
         predict_dataset = predict_dataset.map(
             function=function_kwarg,
             batched=True,