imporve feature_extraction_fn

run_image_captioning_flax.py

@@ -779,17 +779,9 @@ def main():
 
     def feature_extraction_fn(examples):
 
-        pixel_values = []
-
-        for image_file in examples[image_column]:
-            with Image.open(image_file) as image:
-                encoder_inputs = feature_extractor(images=image, return_tensors="np")
-                pixel_values.append(encoder_inputs.pixel_values)
-
-        pixel_values = np.concatenate(pixel_values)
-
-        model_inputs = examples
-        model_inputs['pixel_values'] = pixel_values
+        images = [Image.open(image_file) for image_file in examples[image_column]]
+        encoder_inputs = feature_extractor(images=images, return_tensors="np")
+        model_inputs = {"pixel_values": encoder_inputs.pixel_values}
 
         return model_inputs