# Buggy code
def _construct_index(self, data: pd.DataFrame, predict_mode: bool) -> pd.DataFrame:
    """
    Create index of samples.

    Args:
        data (pd.DataFrame): preprocessed data
        predict_mode (bool): if to create one same per group with prediction length equals ``max_decoder_length``

    Returns:
        pd.DataFrame: index dataframe
    """
    g = data.groupby(self.group_ids, observed=True)

    df_index_first = g["__time_idx__"].transform("nth", 0).to_frame("time_first")
    df_index_last = g["__time_idx__"].transform("nth", -1).to_frame("time_last")
    df_index_diff_to_next = (
        -g["__time_idx__"].diff(-1).fillna(-1).astype(int).to_frame("time_diff_to_next")
    )
    df_index = pd.concat([df_index_first, df_index_last, df_index_diff_to_next], axis=1)
    df_index["index_start"] = np.arange(len(df_index))
    df_index["time"] = data["__time_idx__"]
    df_index["count"] = (df_index["time_last"] - df_index["time_first"]).astype(int) + 1
    group_ids = g.ngroup()
    df_index["group_id"] = group_ids

    min_sequence_length = self.min_prediction_length + self.min_encoder_length
    max_sequence_length = self.max_prediction_length + self.max_encoder_length

    # calculate maximum index to include from current index_start
    max_time = (df_index["time"] + max_sequence_length - 1).clip(
        upper=df_index["count"] + df_index.time_first - 1
    )

    # if there are missing timesteps, we cannot say directly what is the last timestep to include
    # therefore we iterate until it is found
    if (df_index["time_diff_to_next"] != 1).any():
        assert self.allow_missings, (
            "Time difference between steps has been idenfied as larger than 1 - set allow_missings=True"
        )

    df_index["index_end"], missing_sequences = _find_end_indices(
        diffs=df_index.time_diff_to_next.to_numpy(),
        max_lengths=(max_time - df_index.time).to_numpy() + 1,
        min_length=min_sequence_length,
    )
    # add duplicates but mostly with shorter sequence length for start of timeseries
    # while the previous steps have ensured that we start a sequence on every time step, the missing_sequences
    # ensure that there is a sequence that finishes on every timestep
    if len(missing_sequences) > 0:
        shortened_sequences = df_index.iloc[missing_sequences[:, 0]].assign(
            index_end=missing_sequences[:, 1]
        )

        # concatenate shortened sequences
        df_index = pd.concat([df_index, shortened_sequences], axis=0, ignore_index=True)

    # filter out where encode and decode length are not satisfied
    df_index["sequence_length"] = (
        df_index["time"].iloc[df_index["index_end"]].to_numpy() - df_index["time"] + 1
    )

    # filter too short sequences
    df_index = df_index[
        # sequence must be at least of minimal prediction length
        lambda x: (x.sequence_length >= min_sequence_length)
        &
        # prediction must be for after minimal prediction index + length of prediction
        (
            x["sequence_length"] + x["time"]
            >= self.min_prediction_idx + self.min_prediction_length
        )
    ]

    if predict_mode:  # keep longest element per series (i.e. the first element that spans to the end of the series)
        # filter all elements that are longer than the allowed maximum sequence length
        df_index = df_index[
            lambda x: (x["time_last"] - x["time"] + 1 <= max_sequence_length)
            & (x["sequence_length"] >= min_sequence_length)
        ]
        # choose longest sequence
        df_index = df_index.loc[df_index.groupby("group_id").sequence_length.idxmax()]

    assert len(df_index) > 0, "filters should not remove entries"
    # check that all groups/series have at least one entry in the index
    if not group_ids.isin(df_index.group_id).all():
        missing_groups = data.loc[
            ~group_ids.isin(df_index.group_id), self.group_ids
        ].drop_duplicates()
        # decode values
        for name in missing_groups.columns:
            missing_groups[name] = self.transform_values(
                name, missing_groups[name], inverse=True
            )
        raise ValueError(
            f"Min encoder length and/or min prediction length is too large for {len(missing_groups)} series/group. "
            f"First 10 examples: {list(missing_groups.iloc[:10].to_dict(orient='index').values())}"
        )

    return df_index


# 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 _construct_index(self, data: pd.DataFrame, predict_mode: bool) -> pd.DataFrame:
    """
    Create index of samples.

    Args:
        data (pd.DataFrame): preprocessed data
        predict_mode (bool): if to create one same per group with prediction length equals ``max_decoder_length``

    Returns:
        pd.DataFrame: index dataframe
    """
    g = data.groupby(self.group_ids, observed=True)

    df_index_first = g["__time_idx__"].transform("nth", 0).to_frame("time_first")
    df_index_last = g["__time_idx__"].transform("nth", -1).to_frame("time_last")
    df_index_diff_to_next = (
        -g["__time_idx__"].diff(-1).fillna(-1).astype(int).to_frame("time_diff_to_next")
    )
    df_index = pd.concat([df_index_first, df_index_last, df_index_diff_to_next], axis=1)
    df_index["index_start"] = np.arange(len(df_index))
    df_index["time"] = data["__time_idx__"]
    df_index["count"] = (df_index["time_last"] - df_index["time_first"]).astype(int) + 1
    group_ids = g.ngroup()
    df_index["group_id"] = group_ids

    min_sequence_length = self.min_prediction_length + self.min_encoder_length
    max_sequence_length = self.max_prediction_length + self.max_encoder_length

    # calculate maximum index to include from current index_start
    max_time = (df_index["time"] + max_sequence_length - 1).clip(
        upper=df_index["count"] + df_index.time_first - 1
    )

    # if there are missing timesteps, we cannot say directly what is the last timestep to include
    # therefore we iterate until it is found
    if (df_index["time_diff_to_next"] != 1).any():
        assert self.allow_missings, (
            "Time difference between steps has been idenfied as larger than 1 - set allow_missings=True"
        )

    df_index["index_end"], missing_sequences = _find_end_indices(
        diffs=df_index.time_diff_to_next.to_numpy(),
        max_lengths=(max_time - df_index.time).to_numpy() + 1,
        min_length=min_sequence_length,
    )
    # add duplicates but mostly with shorter sequence length for start of timeseries
    # while the previous steps have ensured that we start a sequence on every time step, the missing_sequences
    # ensure that there is a sequence that finishes on every timestep
    if len(missing_sequences) > 0:
        shortened_sequences = df_index.iloc[missing_sequences[:, 0]].assign(
            index_end=missing_sequences[:, 1]
        )

        # concatenate shortened sequences
        df_index = pd.concat([df_index, shortened_sequences], axis=0, ignore_index=True)

    # filter out where encode and decode length are not satisfied
    df_index["sequence_length"] = (
        df_index["time"].iloc[df_index["index_end"]].to_numpy() - df_index["time"] + 1
    )

    # filter too short sequences
    df_index = df_index[
        # sequence must be at least of minimal prediction length
        lambda x: (x.sequence_length >= min_sequence_length)
        &
        # prediction must be for after minimal prediction index + length of prediction
        (
            x["sequence_length"] + x["time"]
            >= self.min_prediction_idx + self.min_prediction_length
        )
    ]

    if predict_mode:  # keep longest element per series (i.e. the first element that spans to the end of the series)
        # filter all elements that are longer than the allowed maximum sequence length
        df_index = df_index[
            lambda x: (x["time_last"] - x["time"] + 1 <= max_sequence_length)
            & (x["sequence_length"] >= min_sequence_length)
        ]
        # choose longest sequence
        df_index = df_index.loc[df_index.groupby("group_id").sequence_length.idxmax()]

    # check that all groups/series have at least one entry in the index
    if not group_ids.isin(df_index.group_id).all():
        missing_groups = data.loc[
            ~group_ids.isin(df_index.group_id), self.group_ids
        ].drop_duplicates()
        # decode values
        for name in missing_groups.columns:
            missing_groups[name] = self.transform_values(
                name, missing_groups[name], inverse=True
            )
        warnings.warn(
            (
                "Min encoder length and/or min_prediction_idx and/or min prediction length is too large for "
                f"{len(missing_groups)} series/groups which therefore are not present in the dataset index. "
                "This means no predictions can be made for those series",
                f"First 10 removed groups: {list(missing_groups.iloc[:10].to_dict(orient='index').values())}",
            ),
            UserWarning,
        )
    assert len(df_index) > 0, "filters should not remove entries"

    return df_index