training protocol fix

NikolayBlagoev · Mar 6, 2024 · 997b49b · 997b49b
1 parent 7e6c696
commit 997b49b
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/trainingprotocol.py b/trainingprotocol.py
@@ -88,7 +88,9 @@ async def start(self):
         await super().start()
         if self.pipeline_rank == 0: 
             try:
-                    batch_idx, (data, target) = next(self.dataloader)
+                    batch_idx, ret = next(self.dataloader)
+                    data = ret['text']
+                    target = ret['text']
             except StopIteration :
                     print("TRAINING COMPLETE")
                     return
@@ -163,7 +165,9 @@ def process_data(self, data:bytes, nodeid, addr):
                     self._apply_grad()
                     # print("\n\n\n\ncalculated")
                     try:
-                        batch_idx, (data, target) = next(self.dataloader)
+                        batch_idx, ret = next(self.dataloader)
+                        data = ret['text']
+                        target = ret['text']
                     except StopIteration :
                         print("TRAINING COMPLETE")
                         return
@@ -208,7 +212,9 @@ def process_data(self, data:bytes, nodeid, addr):
                 self._apply_grad()
                 if self.pipeline_rank == 0:
                     try:
-                        batch_idx, (data, target) = next(self.dataloader)
+                        batch_idx, ret = next(self.dataloader)
+                        data = ret['text']
+                        target = ret['text']
                     except StopIteration :
                         print("TRAINING COMPLETE")
                         return