fixed non contiguous error when using lstm_layers > 1 on gpu (#740)

unit8co · Jan 15, 2022 · ddc84a2 · ddc84a2
1 parent adde52a
commit ddc84a2
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/darts/models/forecasting/tft_model.py b/darts/models/forecasting/tft_model.py
@@ -482,11 +482,15 @@ def forward(self, x) -> Dict[str, torch.Tensor]:
 
         # LSTM
         # calculate initial state
-        input_hidden = self.static_context_hidden_encoder_grn(static_embedding).expand(
-            self.lstm_layers, -1, -1
+        input_hidden = (
+            self.static_context_hidden_encoder_grn(static_embedding)
+            .expand(self.lstm_layers, -1, -1)
+            .contiguous()
         )
-        input_cell = self.static_context_cell_encoder_grn(static_embedding).expand(
-            self.lstm_layers, -1, -1
+        input_cell = (
+            self.static_context_cell_encoder_grn(static_embedding)
+            .expand(self.lstm_layers, -1, -1)
+            .contiguous()
         )
 
         # run local lstm encoder