add bug fix after merging

unit8co · Aug 7, 2023 · 1eb41d8 · 1eb41d8
1 parent 27dd610
commit 1eb41d8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/darts/models/forecasting/transformer_model.py b/darts/models/forecasting/transformer_model.py
@@ -349,6 +349,7 @@ def _prediction_step(self, src: torch.Tensor, tgt: torch.Tensor):
     # Allow teacher forcing
     def training_step(self, train_batch, batch_idx) -> torch.Tensor:
         """performs the training step"""
+        train_batch = list(train_batch)
         future_targets = train_batch[-1]
         train_batch.append(future_targets)
         return super().training_step(train_batch, batch_idx)