Spaces:

flax-community
/

dalle-mini

Running

boris commited on Dec 24, 2021

Commit

15993e3

•

1 Parent(s): 87fed1b

feat: display local TPU's

Files changed (1) hide show

tools/train/train.py CHANGED Viewed

@@ -375,7 +375,7 @@ def main():
         datasets.utils.logging.set_verbosity_error()
         transformers.utils.logging.set_verbosity_error()
-    logger.info(f"TPUs: {jax.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):

         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):