# Buggy code
def _preprocess_data(self, data: pd.DataFrame) -> pd.DataFrame:
    """
    Scale continuous variables, encode categories and set aside target and weight.

    Args:
        data (pd.DataFrame): original data

    Returns:
        pd.DataFrame: pre-processed dataframe
    """

    # encode categoricals
    for name in set(self.categoricals + self.group_ids):
        allow_nans = name in self.dropout_categoricals
        if name in self.variable_groups:  # fit groups
            columns = self.variable_groups[name]
            if name not in self.categorical_encoders:
                self.categorical_encoders[name] = NaNLabelEncoder(
                    add_nan=allow_nans
                ).fit(data[columns].to_numpy().reshape(-1))
            elif self.categorical_encoders[name] is not None:
                try:
                    check_is_fitted(self.categorical_encoders[name])
                except NotFittedError:
                    self.categorical_encoders[name] = self.categorical_encoders[
                        name
                    ].fit(data[columns].to_numpy().reshape(-1))
        else:
            if name not in self.categorical_encoders:
                self.categorical_encoders[name] = NaNLabelEncoder(
                    add_nan=allow_nans
                ).fit(data[name])
            elif self.categorical_encoders[name] is not None:
                try:
                    check_is_fitted(self.categorical_encoders[name])
                except NotFittedError:
                    self.categorical_encoders[name] = self.categorical_encoders[
                        name
                    ].fit(data[name])

    # encode them
    for name in set(self.flat_categoricals + self.group_ids):
        data[name] = self.transform_values(name, data[name], inverse=False)

    # save special variables
    assert "__time_idx__" not in data.columns, (
        "__time_idx__ is a protected column and must not be present in data"
    )
    data["__time_idx__"] = data[self.time_idx]  # save unscaled
    assert "__target__" not in data.columns, (
        "__target__ is a protected column and must not be present in data"
    )
    data["__target__"] = data[self.target]
    if self.weight is not None:
        data["__weight__"] = data[self.weight]

    # train target normalizer
    if self.target_normalizer is not None:
        # fit target normalizer
        try:
            check_is_fitted(self.target_normalizer)
        except NotFittedError:
            if isinstance(self.target_normalizer, EncoderNormalizer):
                self.target_normalizer.fit(data[self.target])
            elif isinstance(self.target_normalizer, GroupNormalizer):
                self.target_normalizer.fit(data[self.target], data)
            else:
                self.target_normalizer.fit(data[self.target])

        # transform target
        if isinstance(self.target_normalizer, EncoderNormalizer):
            # we approximate the scales and target transformation by assuming one
            # transformation over the entire time range but by each group
            common_init_args = [
                name
                for name in inspect.signature(GroupNormalizer).parameters.keys()
                if name in inspect.signature(EncoderNormalizer).parameters.keys()
            ]
            copy_kwargs = {
                name: getattr(self.target_normalizer, name) for name in common_init_args
            }
            normalizer = GroupNormalizer(groups=self.group_ids, **copy_kwargs)
            data[self.target], scales = normalizer.fit_transform(
                data[self.target], data, return_norm=True
            )
        elif isinstance(self.target_normalizer, GroupNormalizer):
            data[self.target], scales = self.target_normalizer.transform(
                data[self.target], data, return_norm=True
            )
        elif isinstance(self.target_normalizer, NaNLabelEncoder):
            data[self.target] = self.target_normalizer.transform(data[self.target])
        else:
            data[self.target], scales = self.target_normalizer.transform(
                data[self.target], return_norm=True
            )

        # add target scales
        if self.add_target_scales:
            for idx, name in enumerate(["center", "scale"]):
                feature_name = f"{self.target}_{name}"
                assert feature_name not in data.columns, (
                    f"{feature_name} is a protected column and must not be present in data"
                )
                data[feature_name] = scales[:, idx].squeeze()
                if feature_name not in self.reals:
                    self.static_reals.append(feature_name)

    if self.target in self.reals:
        self.scalers[self.target] = self.target_normalizer

    # rescale continuous variables apart from target
    for name in self.reals:
        if name not in self.scalers:
            self.scalers[name] = StandardScaler().fit(data[[name]])
        elif self.scalers[name] is not None:
            try:
                check_is_fitted(self.scalers[name])
            except NotFittedError:
                if isinstance(self.scalers[name], GroupNormalizer):
                    self.scalers[name] = self.scalers[name].fit(data[[name]], data)
                else:
                    self.scalers[name] = self.scalers[name].fit(data[[name]])
        if self.scalers[name] is not None and name != self.target:
            data[name] = self.transform_values(
                name, data[name], data=data, inverse=False
            )

    # encode constant values
    self.encoded_constant_fill_strategy = {}
    for name, value in self.constant_fill_strategy.items():
        if name == self.target:
            self.encoded_constant_fill_strategy["__target__"] = value
        self.encoded_constant_fill_strategy[name] = self.transform_values(
            name, np.array([value]), data=data, inverse=False
        )[0]
    return data


# 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 _preprocess_data(self, data: pd.DataFrame) -> pd.DataFrame:
    """
    Scale continuous variables, encode categories and set aside target and weight.

    Args:
        data (pd.DataFrame): original data

    Returns:
        pd.DataFrame: pre-processed dataframe
    """

    # encode categoricals
    for name in set(self.categoricals + self.group_ids):
        allow_nans = name in self.dropout_categoricals
        if name in self.variable_groups:  # fit groups
            columns = self.variable_groups[name]
            if name not in self.categorical_encoders:
                self.categorical_encoders[name] = NaNLabelEncoder(
                    add_nan=allow_nans
                ).fit(data[columns].to_numpy().reshape(-1))
            elif self.categorical_encoders[name] is not None:
                try:
                    check_is_fitted(self.categorical_encoders[name])
                except NotFittedError:
                    self.categorical_encoders[name] = self.categorical_encoders[
                        name
                    ].fit(data[columns].to_numpy().reshape(-1))
        else:
            if name not in self.categorical_encoders:
                self.categorical_encoders[name] = NaNLabelEncoder(
                    add_nan=allow_nans
                ).fit(data[name])
            elif self.categorical_encoders[name] is not None:
                try:
                    check_is_fitted(self.categorical_encoders[name])
                except NotFittedError:
                    self.categorical_encoders[name] = self.categorical_encoders[
                        name
                    ].fit(data[name])

    # encode them
    for name in set(self.flat_categoricals + self.group_ids):
        data[name] = self.transform_values(name, data[name], inverse=False)

    # save special variables
    assert "__time_idx__" not in data.columns, (
        "__time_idx__ is a protected column and must not be present in data"
    )
    data["__time_idx__"] = data[self.time_idx]  # save unscaled
    assert "__target__" not in data.columns, (
        "__target__ is a protected column and must not be present in data"
    )
    data["__target__"] = data[self.target]
    if self.weight is not None:
        data["__weight__"] = data[self.weight]

    # train target normalizer
    if self.target_normalizer is not None:
        # fit target normalizer
        try:
            check_is_fitted(self.target_normalizer)
        except NotFittedError:
            if isinstance(self.target_normalizer, EncoderNormalizer):
                self.target_normalizer.fit(data[self.target])
            elif isinstance(self.target_normalizer, GroupNormalizer):
                self.target_normalizer.fit(data[self.target], data)
            else:
                self.target_normalizer.fit(data[self.target])

        # transform target
        if isinstance(self.target_normalizer, EncoderNormalizer):
            # we approximate the scales and target transformation by assuming one
            # transformation over the entire time range but by each group
            common_init_args = [
                name
                for name in inspect.signature(
                    GroupNormalizer.__init__
                ).parameters.keys()
                if name
                in inspect.signature(EncoderNormalizer.__init__).parameters.keys()
                and name not in ["data", "self"]
            ]
            copy_kwargs = {
                name: getattr(self.target_normalizer, name) for name in common_init_args
            }
            normalizer = GroupNormalizer(groups=self.group_ids, **copy_kwargs)
            data[self.target], scales = normalizer.fit_transform(
                data[self.target], data, return_norm=True
            )
        elif isinstance(self.target_normalizer, GroupNormalizer):
            data[self.target], scales = self.target_normalizer.transform(
                data[self.target], data, return_norm=True
            )
        elif isinstance(self.target_normalizer, NaNLabelEncoder):
            data[self.target] = self.target_normalizer.transform(data[self.target])
        else:
            data[self.target], scales = self.target_normalizer.transform(
                data[self.target], return_norm=True
            )

        # add target scales
        if self.add_target_scales:
            for idx, name in enumerate(["center", "scale"]):
                feature_name = f"{self.target}_{name}"
                assert feature_name not in data.columns, (
                    f"{feature_name} is a protected column and must not be present in data"
                )
                data[feature_name] = scales[:, idx].squeeze()
                if feature_name not in self.reals:
                    self.static_reals.append(feature_name)

    if self.target in self.reals:
        self.scalers[self.target] = self.target_normalizer

    # rescale continuous variables apart from target
    for name in self.reals:
        if name not in self.scalers:
            self.scalers[name] = StandardScaler().fit(data[[name]])
        elif self.scalers[name] is not None:
            try:
                check_is_fitted(self.scalers[name])
            except NotFittedError:
                if isinstance(self.scalers[name], GroupNormalizer):
                    self.scalers[name] = self.scalers[name].fit(data[[name]], data)
                else:
                    self.scalers[name] = self.scalers[name].fit(data[[name]])
        if self.scalers[name] is not None and name != self.target:
            data[name] = self.transform_values(
                name, data[name], data=data, inverse=False
            )

    # encode constant values
    self.encoded_constant_fill_strategy = {}
    for name, value in self.constant_fill_strategy.items():
        if name == self.target:
            self.encoded_constant_fill_strategy["__target__"] = value
        self.encoded_constant_fill_strategy[name] = self.transform_values(
            name, np.array([value]), data=data, inverse=False
        )[0]
    return data