fix

run_image_captioning_flax.py

@@ -929,6 +929,9 @@ def main():
 
     train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
 
+    if training_args.block_size % train_batch_size > 0:
+        raise ValueError(f"`training_args.block_size` needs to be a multiple of the global batch size. Got {training_args.block_size} and {train_batch_size} instead.")
+
     if training_args.do_train:
         steps_per_epoch = len(train_dataset) // train_batch_size
         num_train_examples_per_epoch = steps_per_epoch * train_batch_size