Resume sample counts again

daviswer · web-flow · commit 7d02d2cafbcd · 2025-05-14T10:32:35.000-04:00
diff --git a/fms_fsdp/utils/dataset_utils.py b/fms_fsdp/utils/dataset_utils.py
@@ -1387,7 +1387,7 @@ def __init__(
         self.tokens_seen = [0] * len(self.datasets)
 
         self.current_iterator = -1
-        self.state_params = ["current_iterator"]
+        self.state_params = ["current_iterator", "tokens_seen"]
 
     def setup(self):
         if not self.is_setup:
