Merge pull request #109 from jdb78/fix/stack_variable_lengths_tensors

FIx second stacking variable length tensors occurance
sktime · Oct 17, 2020 · aa4f0d9 · aa4f0d9
2 parents c3992ee + 7930b9a
commit aa4f0d9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_forecasting/models/temporal_fusion_transformer/__init__.py b/pytorch_forecasting/models/temporal_fusion_transformer/__init__.py
@@ -823,7 +823,7 @@ def _log_interpretation(self, outputs, label="train"):
         # log lengths of encoder/decoder
         for type in ["encoder", "decoder"]:
             fig, ax = plt.subplots()
-            lengths = torch.stack([out["interpretation"][f"{type}_length_histogram"] for out in outputs]).sum(0).cpu()
+            lengths = padded_stack([out["interpretation"][f"{type}_length_histogram"] for out in outputs]).sum(0).cpu()
             if type == "decoder":
                 start = 1
             else: