Fix pipline dataloader when batch elements contain tuple (#565)

Co-authored-by: Jeff Rasley <jerasley@microsoft.com> Co-authored-by: Ammar Ahmad Awan <ammar.awan@microsoft.com>
microsoft · Aug 25, 2023 · c69bd1f · c69bd1f
1 parent 0b7a760
commit c69bd1f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
@@ -795,7 +795,8 @@ def _exec_load_micro_batch(self, buffer_id):
             loaded = batch[1]
             if torch.is_tensor(batch[1]):
                 loaded = batch[1].to(self.device)
-            elif isinstance(batch[1], tuple):
+            # XXX: torch 1.6.0 DataLoader will auto convert tuple to list
+            elif isinstance(batch[1], (tuple, list)):
                 loaded = []
                 for x in batch[1]:
                     assert torch.is_tensor(x)