keep in memory for feature extraction

run_image_captioning_flax.py

@@ -871,6 +871,7 @@ def main():
                 remove_columns=[image_column],
                 load_from_cache_file=not data_args.overwrite_cache,
                 features=features,
+                keep_in_memory=True,
                 desc=f"Running feature extraction on {names[split]} dataset".replace("  ", " "),
             )
             _ds = _ds.with_format("numpy")