debug

stanford-crfm · Jan 28, 2024 · 9be382e · 9be382e
1 parent e2c02e4
commit 9be382e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/levanter/data/mixture.py b/src/levanter/data/mixture.py
@@ -36,15 +36,16 @@ def __init__(
         self,
         datasets: Mapping[str, ShardableDataset[T]],
         weights: Dict[str, float],
-        stop_strategy: str = StopStrategy.FIRST_STOP_STRATEGY,
+        # stop_strategy: str = StopStrategy,
+        stop_strategy: str = None,
         key: int | PRNGKeyArray = 0,
     ):
         self.datasets = datasets
         self.weights = MixtureDataset._normalize_weights(weights)
 
         if stop_strategy not in [StopStrategy.FIRST_STOP_STRATEGY, StopStrategy.ALL_STOP_STRATEGY]:
             raise ValueError(f"Stop strategy {stop_strategy} is not supported.")
-
+        print(f"=== Using stop_strategy {stop_strategy} ===")
         self.stop_strategy = stop_strategy
 
         if not isinstance(key, int):