Changed print to logger

run_mlm_flax_stream.py

@@ -700,7 +700,7 @@ if __name__ == "__main__":
 
         # save checkpoint after eval_steps
         if step % training_args.save_steps == 0 and step > 0 and jax.process_index() == 0:
-            print(f"Saving checkpoint at {step + 1} steps")
+            logger.info(f"Saving checkpoint at {step + 1} steps")
             params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
             model.save_pretrained(
                 training_args.output_dir,