Add log info

src/run_clm_flax.py

@@ -64,6 +64,8 @@ from data_utils import (
     normalizer
 )
 
+print(jax.devices())
+
 logger = logging.getLogger(__name__)
 
 # Cache the result
@@ -366,7 +368,7 @@ def main():
     # dataset = dataset.map(normalizer)
     # logger.info(f"Preprocessed dataset kept {len(dataset)} out of {len(raw_dataset)}")
     dataset = raw_dataset
-    
+
     # Load pretrained model and tokenizer
 
     # Distributed training: