fix single dataset error with exhaust with 2b models

2024-08-01 10:37:57 +08:00 · 2024-08-01 10:37:57 +08:00 · ed025abba3
parent 441c79f807
commit ed025abba3
1 changed files with 2 additions and 3 deletions
--- a/FM_9G/fm9g/dragonfly/training_tasks/pretrain_indexed.py
+++ b/FM_9G/fm9g/dragonfly/training_tasks/pretrain_indexed.py
@ -602,8 +602,7 @@ class MixedIndexedDataset(torch.utils.data.IterableDataset):
                idx = np.random.choice(len(self.weights), p=self.weights)

            data = next(self.tasks[idx])
-            if step % self.update_weights_frequency == 0:
-                self.update_weights()
+
            if data is None:
                if self.tasks[idx].allow_repeat:
                    # _runtime_ave = self.tasks[idx].ave_tokens