# Buggy code
def _log_interpretation(self, outputs, label="train"):
    """
    Log interpretation metrics to tensorboard.
    """
    # extract interpretations
    interpretation = {
        name: torch.stack([x["interpretation"][name] for x in outputs]).sum(0)
        for name in outputs[0]["interpretation"].keys()
    }
    # normalize attention with length histogram squared to account for: 1. zeros in attention and
    # 2. higher attention due to less values
    attention_occurances = (
        interpretation["encoder_length_histogram"][1:].flip(0).cumsum(0).float()
    )
    attention_occurances = attention_occurances / attention_occurances.max()
    attention_occurances = torch.cat(
        [
            attention_occurances,
            torch.ones(
                interpretation["attention"].size(0) - attention_occurances.size(0),
                dtype=attention_occurances.dtype,
                device=attention_occurances.device,
            ),
        ],
        dim=0,
    )
    interpretation["attention"] = interpretation[
        "attention"
    ] / attention_occurances.pow(2).clamp(1.0)
    interpretation["attention"] = (
        interpretation["attention"] / interpretation["attention"].sum()
    )

    figs = self.plot_interpretation(interpretation)  # make interpretation figures
    # log to tensorboard
    for name, fig in figs.items():
        self.logger.experiment.add_figure(
            f"{label.capitalize()} {name} importance", fig, global_step=self.global_step
        )

    # 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()
        )
        if type == "decoder":
            start = 1
        else:
            start = 0
        ax.plot(torch.arange(start, start + len(lengths)), lengths)
        ax.set_xlabel(f"{type.capitalize()} length")
        ax.set_ylabel("Number of samples")
        ax.set_title(f"{type.capitalize()} length distribution in {label} epoch")

        self.logger.experiment.add_figure(
            f"{label.capitalize()} {type} length distribution",
            fig,
            global_step=self.global_step,
        )


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-32-14fda4f79b4a> in <module>
1 # Train model
----> 2 trainer.fit(
3     tft,
4     train_dataloader = train_dataloader,
5     val_dataloaders = val_dataloader

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_lightning/trainer/states.py in wrapped_fn(self, *args, **kwargs)
46             if entering is not None:
47                 self.state = entering
---> 48             result = fn(self, *args, **kwargs)
49
50             # The INTERRUPTED state can be set inside the run function. To indicate that run was interrupted

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
1071             self.accelerator_backend = GPUBackend(self)
1072             model = self.accelerator_backend.setup(model)
-> 1073             results = self.accelerator_backend.train(model)
1074
1075         elif self.use_tpu:

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu_backend.py in train(self, model)
49
50     def train(self, model):
---> 51         results = self.trainer.run_pretrain_routine(model)
52         return results
53

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in run_pretrain_routine(self, model)
1237
1238         # CORE TRAINING LOOP
-> 1239         self.train()
1240
1241     def _run_sanity_check(self, ref_model, model):

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in train(self)
392                 # RUN TNG EPOCH
393                 # -----------------
--> 394                 self.run_training_epoch()
395
396                 if self.max_steps and self.max_steps <= self.global_step:

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in run_training_epoch(self)
548
549         # process epoch outputs
--> 550         self.run_training_epoch_end(epoch_output, checkpoint_accumulator, early_stopping_accumulator, num_optimizers)
551
552         # checkpoint callback

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in run_training_epoch_end(self, epoch_output, checkpoint_accumulator, early_stopping_accumulator, num_optimizers)
662             # run training_epoch_end
663             # a list with a result per optimizer index
--> 664             epoch_output = model.training_epoch_end(epoch_output)
665
666             if isinstance(epoch_output, Result):

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_forecasting/models/base_model.py in training_epoch_end(self, outputs)
133
134     def training_epoch_end(self, outputs):
--> 135         log, _ = self.epoch_end(outputs, label="train")
136         return log
137

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_forecasting/models/temporal_fusion_transformer/__init__.py in epoch_end(self, outputs, label)
613         log, out = super().epoch_end(outputs, label=label)
614         if self.log_interval(label == "train") > 0:
--> 615             self._log_interpretation(out, label=label)
616         return log, out
617

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_forecasting/models/temporal_fusion_transformer/__init__.py in _log_interpretation(self, outputs, label)
820         """
821         # extract interpretations
--> 822         interpretation = {
823             name: torch.stack([x["interpretation"][name] for x in outputs]).sum(0)
824             for name in outputs[0]["interpretation"].keys()

~/anaconda3/envs/forecasting/lib/python3.8/site-packages/pytorch_forecasting/models/temporal_fusion_transformer/__init__.py in <dictcomp>(.0)
821         # extract interpretations
822         interpretation = {
--> 823             name: torch.stack([x["interpretation"][name] for x in outputs]).sum(0)
824             for name in outputs[0]["interpretation"].keys()
825         }

RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 6 and 7 in dimension 1 at /tmp/pip-req-build-8yht7tdu/aten/src/THC/generic/THCTensorMath.cu:71

# Fixed code
def _log_interpretation(self, outputs, label="train"):
    """
    Log interpretation metrics to tensorboard.
    """
    # extract interpretations
    interpretation = {
        # use padded_stack because decoder length histogram can be of different length
        name: padded_stack(
            [x["interpretation"][name] for x in outputs], side="right", value=0
        ).sum(0)
        for name in outputs[0]["interpretation"].keys()
    }
    # normalize attention with length histogram squared to account for: 1. zeros in attention and
    # 2. higher attention due to less values
    attention_occurances = (
        interpretation["encoder_length_histogram"][1:].flip(0).cumsum(0).float()
    )
    attention_occurances = attention_occurances / attention_occurances.max()
    attention_occurances = torch.cat(
        [
            attention_occurances,
            torch.ones(
                interpretation["attention"].size(0) - attention_occurances.size(0),
                dtype=attention_occurances.dtype,
                device=attention_occurances.device,
            ),
        ],
        dim=0,
    )
    interpretation["attention"] = interpretation[
        "attention"
    ] / attention_occurances.pow(2).clamp(1.0)
    interpretation["attention"] = (
        interpretation["attention"] / interpretation["attention"].sum()
    )

    figs = self.plot_interpretation(interpretation)  # make interpretation figures
    # log to tensorboard
    for name, fig in figs.items():
        self.logger.experiment.add_figure(
            f"{label.capitalize()} {name} importance", fig, global_step=self.global_step
        )

    # 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()
        )
        if type == "decoder":
            start = 1
        else:
            start = 0
        ax.plot(torch.arange(start, start + len(lengths)), lengths)
        ax.set_xlabel(f"{type.capitalize()} length")
        ax.set_ylabel("Number of samples")
        ax.set_title(f"{type.capitalize()} length distribution in {label} epoch")

        self.logger.experiment.add_figure(
            f"{label.capitalize()} {type} length distribution",
            fig,
            global_step=self.global_step,
        )