add print

stanford-crfm · May 23, 2024 · 96fd096 · 96fd096
1 parent 6a1096a
commit 96fd096
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/levanter/data/loader.py b/src/levanter/data/loader.py
@@ -346,6 +346,8 @@ def _batched(item_iter, size):
 
 def _batched_combine(item_iters, size):
     batch = []
+    if jax.process_index() == 4:
+        logger.info("process 4 enters for loop")
     for items in zip(*item_iters):
         if jax.process_index() == 4:
             logger.info(len(batch))