# Buggy code
def calculate_prediction_actual_by_variable(
    self,
    x: Dict[str, torch.Tensor],
    y_pred: torch.Tensor,
    normalize: bool = True,
    bins: int = 95,
    std: float = 2.0,
) -> Dict[str, Dict[str, torch.Tensor]]:
    """
    Calculate predictions and actuals by variable averaged by ``bins`` bins spanning from ``-std`` to ``+std``

    Args:
        x: input as ``forward()``
        y_pred: predictions obtained by ``self.transform_output(self(x))``
        normalize: if to return normalized averages, i.e. mean or sum of ``y``
        bins: number of bins to calculate
        std: number of standard deviations for standard scaled continuous variables

    Returns:
        dictionary that can be used to plot averages with :py:meth:`~plot_prediction_actual_by_variable`
    """
    support = {}  # histogram
    # averages
    averages_actual = {}
    averages_prediction = {}

    # mask values and transform to log space
    max_encoder_length = x["decoder_lengths"].max()
    mask = self._get_mask(max_encoder_length, x["decoder_lengths"], inverse=True)
    # select valid y values
    y_flat = x["decoder_target"][mask]
    y_pred_flat = y_pred[mask]
    log_y = self.dataset_parameters["target_normalizer"] is not None and getattr(
        self.dataset_parameters["target_normalizer"], "log_scale", False
    )
    if log_y:
        y_flat = torch.log(y_flat + 1e-8)
        y_pred_flat = torch.log(y_pred_flat + 1e-8)

    # real bins
    positive_bins = (bins - 1) // 2

    # if to normalize
    if normalize:
        reduction = "mean"
    else:
        reduction = "sum"
    # continuous variables
    reals = x["decoder_cont"]
    for idx, name in enumerate(self.hparams.x_reals):
        averages_actual[name], support[name] = groupby_apply(
            (reals[..., idx][mask] * positive_bins / std)
            .round()
            .clamp(-positive_bins, positive_bins)
            .long()
            + positive_bins,
            y_flat,
            bins=bins,
            reduction=reduction,
            return_histogram=True,
        )
        averages_prediction[name], _ = groupby_apply(
            (reals[..., idx][mask] * positive_bins / std)
            .round()
            .clamp(-positive_bins, positive_bins)
            .long()
            + positive_bins,
            y_pred_flat,
            bins=bins,
            reduction=reduction,
            return_histogram=True,
        )

    # categorical_variables
    cats = x["decoder_cat"]
    for idx, name in enumerate(
        self.hparams.x_categoricals
    ):  # todo: make it work for grouped categoricals
        reduction = "sum"
        name = self.categorical_groups_mapping.get(name, name)
        averages_actual_cat, support_cat = groupby_apply(
            cats[..., idx][mask],
            y_flat,
            bins=self.hparams.embedding_sizes[name][0],
            reduction=reduction,
            return_histogram=True,
        )
        averages_prediction_cat, _ = groupby_apply(
            cats[..., idx][mask],
            y_pred_flat,
            bins=self.hparams.embedding_sizes[name][0],
            reduction=reduction,
            return_histogram=True,
        )

        # add either to existing calculations or
        if name in averages_actual:
            averages_actual[name] += averages_actual_cat
            support[name] += support_cat
            averages_prediction[name] += averages_prediction_cat
        else:
            averages_actual[name] = averages_actual_cat
            support[name] = support_cat
            averages_prediction[name] = averages_prediction_cat

    if normalize:  # run reduction for categoricals
        for name in self.hparams.embedding_sizes.keys():
            averages_actual[name] /= support[name].clamp(min=1)
            averages_prediction[name] /= support[name].clamp(min=1)

    if log_y:  # reverse log scaling
        for name in support.keys():
            averages_actual[name] = torch.exp(averages_actual[name])
            averages_prediction[name] = torch.exp(averages_prediction[name])

    return {
        "support": support,
        "average": {"actual": averages_actual, "prediction": averages_prediction},
        "std": std,
    }


# 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 calculate_prediction_actual_by_variable(
    self,
    x: Dict[str, torch.Tensor],
    y_pred: torch.Tensor,
    normalize: bool = True,
    bins: int = 95,
    std: float = 2.0,
) -> Dict[str, Dict[str, torch.Tensor]]:
    """
    Calculate predictions and actuals by variable averaged by ``bins`` bins spanning from ``-std`` to ``+std``

    Args:
        x: input as ``forward()``
        y_pred: predictions obtained by ``self.transform_output(self(x, **kwargs))``
        normalize: if to return normalized averages, i.e. mean or sum of ``y``
        bins: number of bins to calculate
        std: number of standard deviations for standard scaled continuous variables

    Returns:
        dictionary that can be used to plot averages with :py:meth:`~plot_prediction_actual_by_variable`
    """
    support = {}  # histogram
    # averages
    averages_actual = {}
    averages_prediction = {}

    # mask values and transform to log space
    max_encoder_length = x["decoder_lengths"].max()
    mask = self._get_mask(max_encoder_length, x["decoder_lengths"], inverse=True)
    # select valid y values
    y_flat = x["decoder_target"][mask]
    y_pred_flat = y_pred[mask]
    log_y = self.dataset_parameters["target_normalizer"] is not None and getattr(
        self.dataset_parameters["target_normalizer"], "log_scale", False
    )
    if log_y:
        y_flat = torch.log(y_flat + 1e-8)
        y_pred_flat = torch.log(y_pred_flat + 1e-8)

    # real bins
    positive_bins = (bins - 1) // 2

    # if to normalize
    if normalize:
        reduction = "mean"
    else:
        reduction = "sum"
    # continuous variables
    reals = x["decoder_cont"]
    for idx, name in enumerate(self.hparams.x_reals):
        averages_actual[name], support[name] = groupby_apply(
            (reals[..., idx][mask] * positive_bins / std)
            .round()
            .clamp(-positive_bins, positive_bins)
            .long()
            + positive_bins,
            y_flat,
            bins=bins,
            reduction=reduction,
            return_histogram=True,
        )
        averages_prediction[name], _ = groupby_apply(
            (reals[..., idx][mask] * positive_bins / std)
            .round()
            .clamp(-positive_bins, positive_bins)
            .long()
            + positive_bins,
            y_pred_flat,
            bins=bins,
            reduction=reduction,
            return_histogram=True,
        )

    # categorical_variables
    cats = x["decoder_cat"]
    for idx, name in enumerate(
        self.hparams.x_categoricals
    ):  # todo: make it work for grouped categoricals
        reduction = "sum"
        name = self.categorical_groups_mapping.get(name, name)
        averages_actual_cat, support_cat = groupby_apply(
            cats[..., idx][mask],
            y_flat,
            bins=self.hparams.embedding_sizes[name][0],
            reduction=reduction,
            return_histogram=True,
        )
        averages_prediction_cat, _ = groupby_apply(
            cats[..., idx][mask],
            y_pred_flat,
            bins=self.hparams.embedding_sizes[name][0],
            reduction=reduction,
            return_histogram=True,
        )

        # add either to existing calculations or
        if name in averages_actual:
            averages_actual[name] += averages_actual_cat
            support[name] += support_cat
            averages_prediction[name] += averages_prediction_cat
        else:
            averages_actual[name] = averages_actual_cat
            support[name] = support_cat
            averages_prediction[name] = averages_prediction_cat

    if normalize:  # run reduction for categoricals
        for name in self.hparams.embedding_sizes.keys():
            averages_actual[name] /= support[name].clamp(min=1)
            averages_prediction[name] /= support[name].clamp(min=1)

    if log_y:  # reverse log scaling
        for name in support.keys():
            averages_actual[name] = torch.exp(averages_actual[name])
            averages_prediction[name] = torch.exp(averages_prediction[name])

    return {
        "support": support,
        "average": {"actual": averages_actual, "prediction": averages_prediction},
        "std": std,
    }