[BUG] Fix issue when training TFT model on mac M1 mps device. element 0 of tensors does not require grad and does not have a grad_fn

pytorch_forecasting/models/base_model.py

@@ -951,6 +951,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,

pytorch_forecasting/models/temporal_fusion_transformer/sub_modules.py

@@ -51,9 +51,16 @@ def __init__(
             self.gate = nn.Sigmoid()
 
     def interpolate(self, x):
-        upsampled = F.interpolate(
-            x.unsqueeze(1), self.output_size, mode="linear", align_corners=True
-        ).squeeze(1)
+        if x.device.type == "mps":
+            x = x.to("cpu")
+            upsampled = F.interpolate(
+                x.unsqueeze(1), self.output_size, mode="linear", align_corners=True
+            ).squeeze(1)
+            upsampled = upsampled.to("mps")
+        else:
+            upsampled = F.interpolate(
+                x.unsqueeze(1), self.output_size, mode="linear", align_corners=True
+            ).squeeze(1)
         if self.trainable:
             upsampled = upsampled * self.gate(self.mask.unsqueeze(0)) * 2.0
         return upsampled
@@ -284,7 +291,8 @@ def __init__(
         prescalers: Dict[str, nn.Linear] = None,
     ):
         """
-        Calculate weights for ``num_inputs`` variables  which are each of size ``input_size``
+        Calculate weights for ``num_inputs`` variables  which are each of size
+        ``input_size``
         """
         super().__init__()