Fix MPS with sequence loss (#2834)

* Add MPS support for list outputs in training eval loop * Add error for invalid state outputs type in trainer * Remove raise ValueError in trainer eval loop --------- Co-authored-by: Daniel King <[email protected]>
mosaicml · Jan 22, 2024 · 1df5557 · 1df5557
1 parent 2fb4519
commit 1df5557
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: