# Buggy code
def predict(
    self,
    data: Union[DataLoader, pd.DataFrame, TimeSeriesDataSet],
    mode: Union[str, Tuple[str, str]] = "prediction",
    return_index: bool = False,
    return_decoder_lengths: bool = False,
    batch_size: int = 64,
    num_workers: int = 0,
    fast_dev_run: bool = False,
    show_progress_bar: bool = False,
    return_x: bool = False,
):
    """
    predict dataloader

    Args:
        dataloader: dataloader, dataframe or dataset
        mode: one of "prediction", "quantiles" or "raw", or tuple ``("raw", output_name)`` where output_name is
            a name in the dictionary returned by ``forward()``
        return_index: if to return the prediction index
        return_decoder_lengths: if to return decoder_lengths
        batch_size: batch size for dataloader - only used if data is not a dataloader is passed
        num_workers: number of workers for dataloader - only used if data is not a dataloader is passed
        fast_dev_run: if to only return results of first batch
        show_progress_bar: if to show progress bar. Defaults to False.
        return_x: if to return network inputs

    Returns:
        output, x, index, decoder_lengths: some elements might not be present depending on what is configured
            to be returned
    """
    # convert to dataloader
    if isinstance(data, pd.DataFrame):
        data = TimeSeriesDataSet.from_parameters(
            self.dataset_parameters, data, predict=True
        )
    if isinstance(data, TimeSeriesDataSet):
        dataloader = data.to_dataloader(
            batch_size=batch_size, train=False, num_workers=num_workers
        )
    else:
        dataloader = data

    # ensure passed dataloader is correct
    assert isinstance(dataloader.dataset, TimeSeriesDataSet), (
        "dataset behind dataloader mut be TimeSeriesDataSet"
    )

    # prepare model
    self.eval()  # no dropout, etc. no gradients

    # run predictions
    output = []
    decode_lenghts = []
    x_list = []
    index = []
    progress_bar = tqdm(
        desc="Predict",
        unit=" batches",
        total=len(dataloader),
        disable=not show_progress_bar,
    )
    with torch.no_grad():
        for x, _ in dataloader:
            # move data to appropriate device
            for name in x.keys():
                if x[name].device != self.device:
                    x[name] = x[name].to(self.device)

            # make prediction
            out = self(x)  # raw output is dictionary
            out["prediction"] = self.transform_output(out)

            lengths = x["decoder_lengths"]
            if return_decoder_lengths:
                decode_lenghts.append(lengths)
            nan_mask = self._get_mask(out["prediction"].size(1), lengths)
            if isinstance(mode, (tuple, list)):
                if mode[0] == "raw":
                    out = out[mode[1]]
                else:
                    raise ValueError(
                        f"If a tuple is specified, the first element must be 'raw' - got {mode[0]} instead"
                    )
            elif mode == "prediction":
                out = self.loss.to_prediction(out["prediction"])
                # mask non-predictions
                out = out.masked_fill(nan_mask, torch.tensor(float("nan")))
            elif mode == "quantiles":
                out = self.loss.to_quantiles(out["prediction"])
                # mask non-predictions
                out = out.masked_fill(
                    nan_mask.unsqueeze(-1), torch.tensor(float("nan"))
                )
            elif mode == "raw":
                pass
            else:
                raise ValueError(f"Unknown mode {mode} - see docs for valid arguments")

            output.append(out)
            if return_x:
                x_list.append(x)
            if return_index:
                index.append(dataloader.dataset.x_to_index(x))
            progress_bar.update()
            if fast_dev_run:
                break

    # concatenate
    if isinstance(mode, (tuple, list)) or mode != "raw":
        output = torch.cat(output, dim=0)
    elif mode == "raw":
        output_cat = {}
        for name in output[0].keys():
            output_cat[name] = torch.cat([out[name] for out in output], dim=0)
        output = output_cat

    # generate output
    if return_x or return_index or return_decoder_lengths:
        output = [output]
    if return_x:
        x_cat = {}
        for name in x_list[0].keys():
            x_cat[name] = torch.cat([x[name] for x in x_list], dim=0)
        x_cat = x_cat
        output.append(x_cat)
    if return_index:
        output.append(pd.concat(index, axis=0, ignore_index=True))
    if return_decoder_lengths:
        output.append(torch.cat(decode_lenghts, dim=0))
    return output


# 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 predict(
    self,
    data: Union[DataLoader, pd.DataFrame, TimeSeriesDataSet],
    mode: Union[str, Tuple[str, str]] = "prediction",
    return_index: bool = False,
    return_decoder_lengths: bool = False,
    batch_size: int = 64,
    num_workers: int = 0,
    fast_dev_run: bool = False,
    show_progress_bar: bool = False,
    return_x: bool = False,
    **kwargs,
):
    """
    predict dataloader

    Args:
        dataloader: dataloader, dataframe or dataset
        mode: one of "prediction", "quantiles" or "raw", or tuple ``("raw", output_name)`` where output_name is
            a name in the dictionary returned by ``forward()``
        return_index: if to return the prediction index
        return_decoder_lengths: if to return decoder_lengths
        batch_size: batch size for dataloader - only used if data is not a dataloader is passed
        num_workers: number of workers for dataloader - only used if data is not a dataloader is passed
        fast_dev_run: if to only return results of first batch
        show_progress_bar: if to show progress bar. Defaults to False.
        return_x: if to return network inputs
        **kwargs: additional arguments to network's forward method

    Returns:
        output, x, index, decoder_lengths: some elements might not be present depending on what is configured
            to be returned
    """
    # convert to dataloader
    if isinstance(data, pd.DataFrame):
        data = TimeSeriesDataSet.from_parameters(
            self.dataset_parameters, data, predict=True
        )
    if isinstance(data, TimeSeriesDataSet):
        dataloader = data.to_dataloader(
            batch_size=batch_size, train=False, num_workers=num_workers
        )
    else:
        dataloader = data

    # ensure passed dataloader is correct
    assert isinstance(dataloader.dataset, TimeSeriesDataSet), (
        "dataset behind dataloader mut be TimeSeriesDataSet"
    )

    # prepare model
    self.eval()  # no dropout, etc. no gradients

    # run predictions
    output = []
    decode_lenghts = []
    x_list = []
    index = []
    progress_bar = tqdm(
        desc="Predict",
        unit=" batches",
        total=len(dataloader),
        disable=not show_progress_bar,
    )
    with torch.no_grad():
        for x, _ in dataloader:
            # move data to appropriate device
            for name in x.keys():
                if x[name].device != self.device:
                    x[name] = x[name].to(self.device)

            # make prediction
            out = self(x, **kwargs)  # raw output is dictionary
            out["prediction"] = self.transform_output(out)

            lengths = x["decoder_lengths"]
            if return_decoder_lengths:
                decode_lenghts.append(lengths)
            nan_mask = self._get_mask(out["prediction"].size(1), lengths)
            if isinstance(mode, (tuple, list)):
                if mode[0] == "raw":
                    out = out[mode[1]]
                else:
                    raise ValueError(
                        f"If a tuple is specified, the first element must be 'raw' - got {mode[0]} instead"
                    )
            elif mode == "prediction":
                out = self.loss.to_prediction(out["prediction"])
                # mask non-predictions
                out = out.masked_fill(nan_mask, torch.tensor(float("nan")))
            elif mode == "quantiles":
                out = self.loss.to_quantiles(out["prediction"])
                # mask non-predictions
                out = out.masked_fill(
                    nan_mask.unsqueeze(-1), torch.tensor(float("nan"))
                )
            elif mode == "raw":
                pass
            else:
                raise ValueError(f"Unknown mode {mode} - see docs for valid arguments")

            output.append(out)
            if return_x:
                x_list.append(x)
            if return_index:
                index.append(dataloader.dataset.x_to_index(x))
            progress_bar.update()
            if fast_dev_run:
                break

    # concatenate
    if isinstance(mode, (tuple, list)) or mode != "raw":
        output = torch.cat(output, dim=0)
    elif mode == "raw":
        output_cat = {}
        for name in output[0].keys():
            output_cat[name] = torch.cat([out[name] for out in output], dim=0)
        output = output_cat

    # generate output
    if return_x or return_index or return_decoder_lengths:
        output = [output]
    if return_x:
        x_cat = {}
        for name in x_list[0].keys():
            x_cat[name] = torch.cat([x[name] for x in x_list], dim=0)
        x_cat = x_cat
        output.append(x_cat)
    if return_index:
        output.append(pd.concat(index, axis=0, ignore_index=True))
    if return_decoder_lengths:
        output.append(torch.cat(decode_lenghts, dim=0))
    return output