fix

run_image_captioning_flax.py

@@ -794,12 +794,15 @@ def main():
 
         if check:
             for k, v in examples.items():
-                examples[k] = v[bools]
+                if k != image_column:
+                    examples[k] = v[bools]
         else:
             assert len(images) == len(examples)
 
         encoder_inputs = feature_extractor(images=images, return_tensors="np")
         model_inputs = {"pixel_values": encoder_inputs.pixel_values}
+        model_inputs.update(examples)
+        model_inputs.pop(image_column)
 
         return model_inputs