FIx second stacking variable length tensors occurance

sktime · Oct 17, 2020 · 7930b9a · 7930b9a
1 parent 0d299c1
commit 7930b9a
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: