load distributed config

ai-hero · May 5, 2024 · d26f2dc · d26f2dc
1 parent 6b00e7b
commit d26f2dc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/aihero/research/finetuning/train.py b/src/aihero/research/finetuning/train.py
@@ -136,7 +136,7 @@ def fetch_dataset(self) -> DatasetDict:
                 raise Exception("Data Abort")
             else:
                 print(f"LOCAL RANK {self.local_rank}: Data ready")
-
+        print(f"LOCAL RANK {self.local_rank}: loading data")
         try:
             splits = {}
             bos_token = self.tokenizer.bos_token
@@ -273,7 +273,7 @@ def fetch_dataset(self) -> DatasetDict:
                 except:  # pylint: disable=bare-except  # noqa: E722
                     print("Unable to create test dataset")
 
-            print("Loaded dataset")
+            print(f"LOCAL RANK {self.local_rank}: data loaded")
             with open(f"{DATASET_DIR}/data_ready.txt", "w") as f:
                 f.write("Data Ready")
             return DatasetDict(splits)