# Buggy code
def plot_prediction(
    self,
    x: Dict[str, torch.Tensor],
    out: Dict[str, torch.Tensor],
    idx: int = 0,
    add_loss_to_title: Union[Metric, bool] = False,
    show_future_observed: bool = True,
    ax=None,
) -> plt.Figure:
    """
    Plot prediction of prediction vs actuals

    Args:
        x: network input
        out: network output
        idx: index of prediction to plot
        add_loss_to_title: if to add loss to title or loss function to calculate. Default to False.
        show_future_observed: if to show actuals for future. Defaults to True.
        ax: matplotlib axes to plot on

    Returns:
        matplotlib figure
    """
    # all true values for y of the first sample in batch
    y_all = torch.cat([x["encoder_target"][idx], x["decoder_target"][idx]])
    if y_all.ndim == 2:  # timesteps, (target, weight), i.e. weight is included
        y_all = y_all[:, 0]
    max_encoder_length = x["encoder_lengths"].max()
    y = torch.cat(
        (
            y_all[: x["encoder_lengths"][idx]],
            y_all[
                max_encoder_length : (max_encoder_length + x["decoder_lengths"][idx])
            ],
        ),
    )
    # get predictions
    y_pred = out["prediction"].detach().cpu()
    y_hat = y_pred[idx, : x["decoder_lengths"][idx]]

    # move to cpu
    y = y.detach().cpu()
    # create figure
    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = ax.get_figure()
    n_pred = y_hat.shape[0]
    x_obs = np.arange(-(y.shape[0] - n_pred), 0)
    x_pred = np.arange(n_pred)
    prop_cycle = iter(plt.rcParams["axes.prop_cycle"])
    obs_color = next(prop_cycle)["color"]
    pred_color = next(prop_cycle)["color"]
    # plot observed history
    if len(x_obs) > 0:
        if len(x_obs) > 1:
            plotter = ax.plot
        else:
            plotter = ax.scatter
        plotter(x_obs, y[:-n_pred], label="observed", c=obs_color)
    if len(x_pred) > 1:
        plotter = ax.plot
    else:
        plotter = ax.scatter

    # plot observed prediction
    if show_future_observed:
        plotter(x_pred, y[-n_pred:], label=None, c=obs_color)

    # plot prediction
    plotter(
        x_pred,
        self.loss.to_prediction(y_hat.unsqueeze(0))[0],
        label="predicted",
        c=pred_color,
    )

    # plot predicted quantiles
    y_quantiles = self.loss.to_quantiles(y_hat.unsqueeze(0))[0]
    plotter(x_pred, y_quantiles[:, y_quantiles.shape[1] // 2], c=pred_color, alpha=0.15)
    for i in range(y_quantiles.shape[1] // 2):
        if len(x_pred) > 1:
            ax.fill_between(
                x_pred,
                y_quantiles[:, i],
                y_quantiles[:, -i - 1],
                alpha=0.15,
                fc=pred_color,
            )
        else:
            quantiles = torch.tensor([[y_quantiles[0, i]], [y_quantiles[0, -i - 1]]])
            ax.errorbar(
                x_pred,
                y[[-n_pred]],
                yerr=quantiles - y[-n_pred],
                c=pred_color,
                capsize=1.0,
            )
    if add_loss_to_title:
        if isinstance(add_loss_to_title, bool):
            loss = self.loss
        else:
            loss = add_loss_to_title
            loss.quantiles = self.loss.quantiles
        if isinstance(loss, MASE):
            loss_value = loss(y_hat[None], y[-n_pred:][None], y[:n_pred][None])
        else:
            loss_value = loss(y_hat[None], y[-n_pred:][None])
        ax.set_title(f"Loss {loss_value:.3g}")
    ax.set_xlabel("Time index")
    fig.legend()
    return fig


# Traceback (RuntimeError)
RuntimeError                              Traceback (most recent call last)
<ipython-input-11-a92b5627800b> in <module>
5     val_dataloaders=val_dataloader,
6     max_lr=10.0,
----> 7     min_lr=1e-6,
8 )
9

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/tuner/tuning.py in lr_find(self, model, train_dataloader, val_dataloaders, min_lr, max_lr, num_training, mode, early_stop_threshold, datamodule)
128             mode,
129             early_stop_threshold,
--> 130             datamodule,
131         )
132

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/tuner/lr_finder.py in lr_find(trainer, model, train_dataloader, val_dataloaders, min_lr, max_lr, num_training, mode, early_stop_threshold, datamodule)
173                 train_dataloader=train_dataloader,
174                 val_dataloaders=val_dataloaders,
--> 175                 datamodule=datamodule)
176
177     # Prompt if we stopped early

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
437         self.call_hook('on_fit_start')
438
--> 439         results = self.accelerator_backend.train()
440         self.accelerator_backend.teardown()
441

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py in train(self)
52
53         # train or test
---> 54         results = self.train_or_test()
55         return results
56

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py in train_or_test(self)
64             results = self.trainer.run_test()
65         else:
---> 66             results = self.trainer.train()
67         return results
68

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in train(self)
459
460     def train(self):
--> 461         self.run_sanity_check(self.get_model())
462
463         # enable train mode

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in run_sanity_check(self, ref_model)
645
646             # run eval step
--> 647             _, eval_results = self.run_evaluation(test_mode=False, max_batches=self.num_sanity_val_batches)
648
649             # allow no returns from eval

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in run_evaluation(self, test_mode, max_batches)
565
566                 # lightning module methods
--> 567                 output = self.evaluation_loop.evaluation_step(test_mode, batch, batch_idx, dataloader_idx)
568                 output = self.evaluation_loop.evaluation_step_end(output)
569

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py in evaluation_step(self, test_mode, batch, batch_idx, dataloader_idx)
169             output = self.trainer.accelerator_backend.test_step(args)
170         else:
--> 171             output = self.trainer.accelerator_backend.validation_step(args)
172
173         # track batch size for weighted average

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py in validation_step(self, args)
76                 output = self.__validation_step(args)
77         else:
---> 78             output = self.__validation_step(args)
79
80         return output

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py in __validation_step(self, args)
84         batch = self.to_device(batch)
85         args[0] = batch
---> 86         output = self.trainer.model.validation_step(*args)
87         return output
88

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_forecasting/models/base_model.py in validation_step(self, batch, batch_idx)
138     def validation_step(self, batch, batch_idx):
139         x, y = batch
--> 140         log, _ = self.step(x, y, batch_idx, label="val")  # log loss
141         self.log("val_loss", log["loss"], on_step=False, on_epoch=True, prog_bar=True)
142         return log

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_forecasting/models/temporal_fusion_transformer/__init__.py in step(self, x, y, batch_idx, label)
566         """
567         # extract data and run model
--> 568         log, out = super().step(x, y, batch_idx, label=label)
569         # calculate interpretations etc for latter logging
570         if self.log_interval(label == "train") > 0:

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_forecasting/models/base_model.py in step(self, x, y, batch_idx, label)
194             loss = loss * (1 + monotinicity_loss)
195         else:
--> 196             out = self(x)
197             out["prediction"] = self.transform_output(out)
198

~/repo/emart-promo/env/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
725             result = self._slow_forward(*input, **kwargs)
726         else:
--> 727             result = self.forward(*input, **kwargs)
728         for hook in itertools.chain(
729                 _global_forward_hooks.values(),

~/repo/emart-promo/env/lib/python3.7/site-packages/pytorch_forecasting/models/temporal_fusion_transformer/__init__.py in forward(self, x)
489         encoder_output, (hidden, cell) = self.lstm_encoder(
490             rnn.pack_padded_sequence(
--> 491                 embeddings_varying_encoder, lstm_encoder_lengths, enforce_sorted=False, batch_first=True
492             ),
493             (input_hidden, input_cell),

~/repo/emart-promo/env/lib/python3.7/site-packages/torch/nn/utils/rnn.py in pack_padded_sequence(input, lengths, batch_first, enforce_sorted)
242
243     data, batch_sizes = \
--> 244         _VF._pack_padded_sequence(input, lengths, batch_first)
245     return _packed_sequence_init(data, batch_sizes, sorted_indices, None)
246

RuntimeError: 'lengths' argument should be a 1D CPU int64 tensor, but got 1D cuda:0 Long tensor

# Fixed code
def plot_prediction(
    self,
    x: Dict[str, torch.Tensor],
    out: Dict[str, torch.Tensor],
    idx: int = 0,
    add_loss_to_title: Union[Metric, torch.Tensor, bool] = False,
    show_future_observed: bool = True,
    ax=None,
) -> plt.Figure:
    """
    Plot prediction of prediction vs actuals

    Args:
        x: network input
        out: network output
        idx: index of prediction to plot
        add_loss_to_title: if to add loss to title or loss function to calculate. Can be either metrics,
            bool indicating if to use loss metric or tensor which contains losses for all samples. Default to False.
        show_future_observed: if to show actuals for future. Defaults to True.
        ax: matplotlib axes to plot on

    Returns:
        matplotlib figure
    """
    # all true values for y of the first sample in batch
    y_all = torch.cat([x["encoder_target"][idx], x["decoder_target"][idx]])
    if y_all.ndim == 2:  # timesteps, (target, weight), i.e. weight is included
        y_all = y_all[:, 0]
    max_encoder_length = x["encoder_lengths"].max()
    y = torch.cat(
        (
            y_all[: x["encoder_lengths"][idx]],
            y_all[
                max_encoder_length : (max_encoder_length + x["decoder_lengths"][idx])
            ],
        ),
    )
    # get predictions
    y_pred = out["prediction"].detach().cpu()
    y_hat = y_pred[idx, : x["decoder_lengths"][idx]]

    # move to cpu
    y = y.detach().cpu()
    # create figure
    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = ax.get_figure()
    n_pred = y_hat.shape[0]
    x_obs = np.arange(-(y.shape[0] - n_pred), 0)
    x_pred = np.arange(n_pred)
    prop_cycle = iter(plt.rcParams["axes.prop_cycle"])
    obs_color = next(prop_cycle)["color"]
    pred_color = next(prop_cycle)["color"]
    # plot observed history
    if len(x_obs) > 0:
        if len(x_obs) > 1:
            plotter = ax.plot
        else:
            plotter = ax.scatter
        plotter(x_obs, y[:-n_pred], label="observed", c=obs_color)
    if len(x_pred) > 1:
        plotter = ax.plot
    else:
        plotter = ax.scatter

    # plot observed prediction
    if show_future_observed:
        plotter(x_pred, y[-n_pred:], label=None, c=obs_color)

    # plot prediction
    plotter(
        x_pred,
        self.loss.to_prediction(y_hat.unsqueeze(0))[0],
        label="predicted",
        c=pred_color,
    )

    # plot predicted quantiles
    y_quantiles = self.loss.to_quantiles(y_hat.unsqueeze(0))[0]
    plotter(x_pred, y_quantiles[:, y_quantiles.shape[1] // 2], c=pred_color, alpha=0.15)
    for i in range(y_quantiles.shape[1] // 2):
        if len(x_pred) > 1:
            ax.fill_between(
                x_pred,
                y_quantiles[:, i],
                y_quantiles[:, -i - 1],
                alpha=0.15,
                fc=pred_color,
            )
        else:
            quantiles = torch.tensor([[y_quantiles[0, i]], [y_quantiles[0, -i - 1]]])
            ax.errorbar(
                x_pred,
                y[[-n_pred]],
                yerr=quantiles - y[-n_pred],
                c=pred_color,
                capsize=1.0,
            )
    if add_loss_to_title is not False:
        if isinstance(add_loss_to_title, bool):
            loss = self.loss
        elif isinstance(add_loss_to_title, torch.Tensor):
            loss = add_loss_to_title.detach()[idx].item()
        elif isinstance(add_loss_to_title, Metric):
            loss = add_loss_to_title
            loss.quantiles = self.loss.quantiles
        else:
            raise ValueError(f"add_loss_to_title '{add_loss_to_title}'' is unkown")
        if isinstance(loss, MASE):
            loss_value = loss(y_hat[None], y[-n_pred:][None], y[:n_pred][None])
        elif isinstance(loss, Metric):
            loss_value = loss(y_hat[None], y[-n_pred:][None])
        else:
            loss_value = loss
        ax.set_title(f"Loss {loss_value:.3g}")
    ax.set_xlabel("Time index")
    fig.legend()
    return fig