Fix

EasyLM/data.py

@@ -227,7 +227,7 @@ class HuggingfaceDataset(object):
     def get_state_dict(self):
         return dict(
             config=self.config,
-            dataset_loc=self._index if self._train_epochs < 1 else 0,
+            dataset_loc=self._index,
             total_tokens=self._total_tokens,
             epochs=self._train_epochs,
         )