add some comments

InternLM · Sep 25, 2024 · f41cf97 · f41cf97
1 parent 0bc4c44
commit f41cf97
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/internlm/data/megatron/batch_sampler.py b/internlm/data/megatron/batch_sampler.py
@@ -57,6 +57,6 @@ def __iter__(self):
             start_idx, end_idx = self.get_start_end_idx()
             yield batch[start_idx:end_idx]
 
-    # TODO: implement copy method that compatible with InternEvo trainstate
+    # TODO: implement copy method that compatible with InternEvo trainstate ckpt save and load.
     def copy(self):
         return copy.deepcopy(self)
diff --git a/internlm/data/mocked/batch_sampler.py b/internlm/data/mocked/batch_sampler.py
@@ -19,6 +19,6 @@ def __iter__(self):
     def __len__(self):
         return (len(self.train_ds) + self.micro_num - 1) // self.micro_num
 
-    # TODO: implement copy method that compatible with InternEvo trainstate
+    # TODO: implement copy method that compatible with InternEvo trainstate ckpt save and load.
     def copy(self):
         return copy.deepcopy(self)