# Buggy code
def forward(self, x: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """
    input dimensions: n_samples x time x variables
    """
    encoder_lengths = x["encoder_lengths"]
    decoder_lengths = x["decoder_lengths"]
    x_cat = torch.cat(
        [x["encoder_cat"], x["decoder_cat"]], dim=1
    )  # concatenate in time dimension
    x_cont = torch.cat(
        [x["encoder_cont"], x["decoder_cont"]], dim=1
    )  # concatenate in time dimension
    timesteps = x_cont.size(1)  # encode + decode length
    max_encoder_length = int(encoder_lengths.max())
    input_vectors = self.input_embeddings(x_cat)
    input_vectors.update(
        {
            name: x_cont[..., idx].unsqueeze(-1)
            for idx, name in enumerate(self.hparams.x_reals)
        }
    )

    # Embedding and variable selection
    if len(self.hparams.static_categoricals + self.hparams.static_reals) > 0:
        # static embeddings will be constant over entire batch
        static_embedding = {
            name: input_vectors[name][:, 0]
            for name in self.hparams.static_categoricals + self.hparams.static_reals
        }
        static_embedding, static_variable_selection = self.static_variable_selection(
            static_embedding
        )
    else:
        static_embedding = torch.zeros(
            (x_cont.size(0), self.hparams.hidden_size),
            dtype=self.dtype,
            device=self.device,
        )
        static_variable_selection = torch.zeros(
            (x_cont.size(0), 0), dtype=self.dtype, device=self.device
        )

    static_context_variable_selection = self.expand_static_context(
        self.static_context_variable_selection(static_embedding), timesteps
    )

    embeddings_varying_encoder = {
        name: input_vectors[name][:, :max_encoder_length]
        for name in self.hparams.time_varying_categoricals_encoder
        + self.hparams.time_varying_reals_encoder
    }
    embeddings_varying_encoder, encoder_sparse_weights = (
        self.encoder_variable_selection(
            embeddings_varying_encoder,
            static_context_variable_selection[:, :max_encoder_length],
        )
    )

    embeddings_varying_decoder = {
        name: input_vectors[name][:, max_encoder_length:]  # select decoder
        for name in self.hparams.time_varying_categoricals_decoder
        + self.hparams.time_varying_reals_decoder
    }
    embeddings_varying_decoder, decoder_sparse_weights = (
        self.decoder_variable_selection(
            embeddings_varying_decoder,
            static_context_variable_selection[:, max_encoder_length:],
        )
    )

    # LSTM
    # run lstm at least once, i.e. encode length has to be > 0
    lstm_encoder_lengths = encoder_lengths.where(
        encoder_lengths > 0, torch.ones_like(encoder_lengths)
    )
    # calculate initial state
    input_hidden = self.static_context_initial_hidden_lstm(static_embedding).expand(
        self.hparams.lstm_layers, -1, -1
    )
    input_cell = self.static_context_initial_cell_lstm(static_embedding).expand(
        self.hparams.lstm_layers, -1, -1
    )

    # # run local encoder
    encoder_output, (hidden, cell) = self.lstm_encoder(
        rnn.pack_padded_sequence(
            embeddings_varying_encoder,
            lstm_encoder_lengths,
            enforce_sorted=False,
            batch_first=True,
        ),
        (input_hidden, input_cell),
    )
    encoder_output, _ = rnn.pad_packed_sequence(encoder_output, batch_first=True)
    # replace hidden cell with initial input if encoder_length is zero to determine correct initial state
    no_encoding = (encoder_lengths == 0)[
        None, :, None
    ]  # shape: n_lstm_layers x batch_size x hidden_size
    hidden = hidden.masked_scatter(no_encoding, input_hidden)
    cell = cell.masked_scatter(no_encoding, input_cell)

    # run local decoder
    decoder_output, _ = self.lstm_decoder(
        rnn.pack_padded_sequence(
            embeddings_varying_decoder,
            decoder_lengths,
            enforce_sorted=False,
            batch_first=True,
        ),
        (hidden, cell),
    )

    decoder_output, _ = rnn.pad_packed_sequence(decoder_output, batch_first=True)

    # skip connection over lstm
    lstm_output_encoder = self.post_lstm_gate_encoder(encoder_output)
    lstm_output_encoder = self.post_lstm_add_norm_encoder(
        lstm_output_encoder, embeddings_varying_encoder
    )

    lstm_output_decoder = self.post_lstm_gate_decoder(decoder_output)
    lstm_output_decoder = self.post_lstm_add_norm_decoder(
        lstm_output_decoder, embeddings_varying_decoder
    )

    lstm_output = torch.cat([lstm_output_encoder, lstm_output_decoder], dim=1)

    # static enrichment
    static_context_enrichment = self.static_context_enrichment(static_embedding)
    attn_input = self.static_enrichment(
        lstm_output, self.expand_static_context(static_context_enrichment, timesteps)
    )

    # Attention
    attn_output, attn_output_weights = self.multihead_attn(
        q=attn_input[:, max_encoder_length:],  # query only for predictions
        k=attn_input,
        v=attn_input,
        mask=self.get_attention_mask(
            encoder_lengths=encoder_lengths,
            decoder_length=timesteps - max_encoder_length,
        ),
    )

    # skip connection over attention
    attn_output = self.post_attn_gate_norm(
        attn_output, attn_input[:, max_encoder_length:]
    )

    output = self.pos_wise_ff(attn_output)

    # skip connection over temporal fusion decoder (not LSTM decoder despite the LSTM output contains
    # a skip from the variable selection network)
    output = self.pre_output_gate_norm(output, lstm_output[:, max_encoder_length:])
    output = self.output_layer(output)

    return dict(
        prediction=output,
        attention=attn_output_weights,
        static_variables=static_variable_selection,
        encoder_variables=encoder_sparse_weights,
        decoder_variables=decoder_sparse_weights,
        decoder_lengths=decoder_lengths,
        encoder_lengths=encoder_lengths,
        groups=x["groups"],
        decoder_time_idx=x["decoder_time_idx"],
        target_scale=x["target_scale"],
    )


# 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 forward(self, x: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """
    input dimensions: n_samples x time x variables
    """
    encoder_lengths = x["encoder_lengths"]
    decoder_lengths = x["decoder_lengths"]
    x_cat = torch.cat(
        [x["encoder_cat"], x["decoder_cat"]], dim=1
    )  # concatenate in time dimension
    x_cont = torch.cat(
        [x["encoder_cont"], x["decoder_cont"]], dim=1
    )  # concatenate in time dimension
    timesteps = x_cont.size(1)  # encode + decode length
    max_encoder_length = int(encoder_lengths.max())
    input_vectors = self.input_embeddings(x_cat)
    input_vectors.update(
        {
            name: x_cont[..., idx].unsqueeze(-1)
            for idx, name in enumerate(self.hparams.x_reals)
            if name in self.reals
        }
    )

    # Embedding and variable selection
    if len(self.static_variables) > 0:
        # static embeddings will be constant over entire batch
        static_embedding = {
            name: input_vectors[name][:, 0] for name in self.static_variables
        }
        static_embedding, static_variable_selection = self.static_variable_selection(
            static_embedding
        )
    else:
        static_embedding = torch.zeros(
            (x_cont.size(0), self.hparams.hidden_size),
            dtype=self.dtype,
            device=self.device,
        )
        static_variable_selection = torch.zeros(
            (x_cont.size(0), 0), dtype=self.dtype, device=self.device
        )

    static_context_variable_selection = self.expand_static_context(
        self.static_context_variable_selection(static_embedding), timesteps
    )

    embeddings_varying_encoder = {
        name: input_vectors[name][:, :max_encoder_length]
        for name in self.encoder_variables
    }
    embeddings_varying_encoder, encoder_sparse_weights = (
        self.encoder_variable_selection(
            embeddings_varying_encoder,
            static_context_variable_selection[:, :max_encoder_length],
        )
    )

    embeddings_varying_decoder = {
        name: input_vectors[name][:, max_encoder_length:]
        for name in self.decoder_variables  # select decoder
    }
    embeddings_varying_decoder, decoder_sparse_weights = (
        self.decoder_variable_selection(
            embeddings_varying_decoder,
            static_context_variable_selection[:, max_encoder_length:],
        )
    )

    # LSTM
    # run lstm at least once, i.e. encode length has to be > 0
    lstm_encoder_lengths = encoder_lengths.where(
        encoder_lengths > 0, torch.ones_like(encoder_lengths)
    )
    # calculate initial state
    input_hidden = self.static_context_initial_hidden_lstm(static_embedding).expand(
        self.hparams.lstm_layers, -1, -1
    )
    input_cell = self.static_context_initial_cell_lstm(static_embedding).expand(
        self.hparams.lstm_layers, -1, -1
    )

    # # run local encoder
    encoder_output, (hidden, cell) = self.lstm_encoder(
        rnn.pack_padded_sequence(
            embeddings_varying_encoder,
            lstm_encoder_lengths.cpu(),
            enforce_sorted=False,
            batch_first=True,
        ),
        (input_hidden, input_cell),
    )
    encoder_output, _ = rnn.pad_packed_sequence(encoder_output, batch_first=True)
    # replace hidden cell with initial input if encoder_length is zero to determine correct initial state
    no_encoding = (encoder_lengths == 0)[
        None, :, None
    ]  # shape: n_lstm_layers x batch_size x hidden_size
    hidden = hidden.masked_scatter(no_encoding, input_hidden)
    cell = cell.masked_scatter(no_encoding, input_cell)

    # run local decoder
    decoder_output, _ = self.lstm_decoder(
        rnn.pack_padded_sequence(
            embeddings_varying_decoder,
            decoder_lengths.cpu(),
            enforce_sorted=False,
            batch_first=True,
        ),
        (hidden, cell),
    )

    decoder_output, _ = rnn.pad_packed_sequence(decoder_output, batch_first=True)

    # skip connection over lstm
    lstm_output_encoder = self.post_lstm_gate_encoder(encoder_output)
    lstm_output_encoder = self.post_lstm_add_norm_encoder(
        lstm_output_encoder, embeddings_varying_encoder
    )

    lstm_output_decoder = self.post_lstm_gate_decoder(decoder_output)
    lstm_output_decoder = self.post_lstm_add_norm_decoder(
        lstm_output_decoder, embeddings_varying_decoder
    )

    lstm_output = torch.cat([lstm_output_encoder, lstm_output_decoder], dim=1)

    # static enrichment
    static_context_enrichment = self.static_context_enrichment(static_embedding)
    attn_input = self.static_enrichment(
        lstm_output, self.expand_static_context(static_context_enrichment, timesteps)
    )

    # Attention
    attn_output, attn_output_weights = self.multihead_attn(
        q=attn_input[:, max_encoder_length:],  # query only for predictions
        k=attn_input,
        v=attn_input,
        mask=self.get_attention_mask(
            encoder_lengths=encoder_lengths,
            decoder_length=timesteps - max_encoder_length,
        ),
    )

    # skip connection over attention
    attn_output = self.post_attn_gate_norm(
        attn_output, attn_input[:, max_encoder_length:]
    )

    output = self.pos_wise_ff(attn_output)

    # skip connection over temporal fusion decoder (not LSTM decoder despite the LSTM output contains
    # a skip from the variable selection network)
    output = self.pre_output_gate_norm(output, lstm_output[:, max_encoder_length:])
    output = self.output_layer(output)

    return dict(
        prediction=output,
        attention=attn_output_weights,
        static_variables=static_variable_selection,
        encoder_variables=encoder_sparse_weights,
        decoder_variables=decoder_sparse_weights,
        decoder_lengths=decoder_lengths,
        encoder_lengths=encoder_lengths,
        groups=x["groups"],
        decoder_time_idx=x["decoder_time_idx"],
        target_scale=x["target_scale"],
    )