Add MPS support for list outputs in training eval loop

mosaicml · Jan 18, 2024 · a2de08f · a2de08f
1 parent 09dbe18
commit a2de08f
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
@@ -2983,6 +2983,13 @@ def _eval_loop(
                                                 outputs[k] = v.cpu()
                                             else:
                                                 outputs[k] = v
+                                    elif isinstance(self.state.outputs, Sequence):
+                                        outputs = []
+                                        for v in self.state.outputs:
+                                            if isinstance(v, torch.Tensor):
+                                                outputs.append(v.cpu())
+                                            else:
+                                                outputs.append(v)
                                     else:
                                         outputs = self.state.outputs.cpu()
                                 else: