fix issue when there is mps fallback enabled during training

sktime · Dec 9, 2024 · a7f6491 · a7f6491
1 parent d57b0bb
commit a7f6491
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/pytorch_forecasting/models/base_model.py b/pytorch_forecasting/models/base_model.py
@@ -828,6 +828,9 @@ def step(
                     loss = self.loss(prediction, y)
             else:
                 loss = None
+        # ensure that loss has require_grad
+        if loss is not None and loss.device.type == "mps":
+            loss.requires_grad_(True)
         self.log(
             f"{self.current_stage}_loss",
             loss,