add suggested fix from david

stanford-crfm · Oct 23, 2024 · 1706803 · 1706803
1 parent fd18cae
commit 1706803
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/levanter/callbacks.py b/src/levanter/callbacks.py
@@ -18,8 +18,7 @@
 from tqdm_loggable.auto import tqdm
 
 import levanter.tracker
-from levanter.data import DataLoader
-from levanter.data.text import TokenSeqEpochDataset
+from levanter.data import DataLoader, AsyncDataset
 from levanter.logging import save_xla_dumps_to_wandb
 from levanter.tracker.helpers import log_optimizer_hyperparams
 from levanter.tracker.wandb import WandbConfig
@@ -55,7 +54,7 @@ def log_epoch(step_info: StepInfo):
     return log_epoch
 
 
-def get_total_dataset_tokens(ds: TokenSeqEpochDataset, seq_length: int):
+def get_total_dataset_tokens(ds: AsyncDataset, seq_length: int):
     if not ds.is_finite():
        raise ValueError("Epochs don't make sense with an infinite dataset.")