feat: load data first

tools/train/train.py

@@ -375,9 +375,6 @@ def main():
         datasets.utils.logging.set_verbosity_error()
         transformers.utils.logging.set_verbosity_error()
 
-    logger.info(f"Local TPUs: {jax.local_device_count()}")
-    assert jax.local_device_count() == 8, "TPUs in use, please check running processes"
-
     # Set the verbosity to info of the Transformers logger (on main process only):
     logger.info(f"Training/evaluation parameters {training_args}")
 
@@ -388,6 +385,9 @@ def main():
         do_eval=training_args.do_eval,
     )
 
+    logger.info(f"Local TPUs: {jax.local_device_count()}")
+    assert jax.local_device_count() == 8, "TPUs in use, please check running processes"
+
     # Set up wandb run
     if jax.process_index() == 0:
         wandb.init(