# Buggy code
def __init__(
    self,
    hidden_size: int = 16,
    lstm_layers: int = 1,
    dropout: float = 0.1,
    output_size: int = 7,
    loss: MultiHorizonMetric = None,
    attention_head_size: int = 4,
    max_encoder_length: int = 10,
    static_categoricals: List[str] = [],
    static_reals: List[str] = [],
    time_varying_categoricals_encoder: List[str] = [],
    time_varying_categoricals_decoder: List[str] = [],
    categorical_groups: Dict[str, List[str]] = {},
    time_varying_reals_encoder: List[str] = [],
    time_varying_reals_decoder: List[str] = [],
    x_reals: List[str] = [],
    x_categoricals: List[str] = [],
    hidden_continuous_size: int = 8,
    hidden_continuous_sizes: Dict[str, int] = {},
    embedding_sizes: Dict[str, Tuple[int, int]] = {},
    embedding_paddings: List[str] = [],
    embedding_labels: Dict[str, np.ndarray] = {},
    learning_rate: float = 1e-3,
    log_interval: Union[int, float] = -1,
    log_val_interval: Union[int, float] = None,
    log_gradient_flow: bool = False,
    reduce_on_plateau_patience: int = 1000,
    monotone_constaints: Dict[str, int] = {},
    share_single_variable_networks: bool = False,
    logging_metrics: nn.ModuleList = None,
    **kwargs,
):
    """
    Temporal Fusion Transformer for forecasting timeseries - use its :py:meth:`~from_dataset` method if possible.

    Implementation of the article
    `Temporal Fusion Transformers for Interpretable Multi-horizon Time Series
    Forecasting <https://arxiv.org/pdf/1912.09363.pdf>`_. The network outperforms DeepAR by Amazon by 36-69%
    in benchmarks.

    Enhancements compared to the original implementation (apart from capabilities added through base model
    such as monotone constraints):

    * static variables can be continuous
    * multiple categorical variables can be summarized with an EmbeddingBag
    * variable encoder and decoder length by sample
    * categorical embeddings are not transformed by variable selection network (because it is a redundant operation)
    * variable dimension in variable selection network are scaled up via linear interpolation to reduce
      number of parameters
    * non-linear variable processing in variable selection network can be shared among decoder and encoder
      (not shared by default)

    Tune its hyperparameters with
    :py:func:`~pytorch_forecasting.models.temporal_fusion_transformer.tuning.optimize_hyperparameters`.

    Args:

        hidden_size: hidden size of network which is its main hyperparameter and can range from 8 to 512
        lstm_layers: number of LSTM layers (2 is mostly optimal)
        dropout: dropout rate
        output_size: number of outputs (e.g. number of quantiles for QuantileLoss)
        loss: loss function taking prediction and targets
        attention_head_size: number of attention heads (4 is a good default)
        max_encoder_length: length to encode (can be far longer than the decoder length but does not have to be)
        static_categoricals: integer of positions of static categorical variables
        static_reals: integer of positions of static continuous variables
        time_varying_categoricals_encoder: integer of positions of categorical variables for encoder
        time_varying_categoricals_decoder: integer of positions of categorical variables for decoder
        time_varying_reals_encoder: integer of positions of continuous variables for encoder
        time_varying_reals_decoder: integer of positions of continuous variables for decoder
        categorical_groups: dictionary where values
            are list of categorical variables that are forming together a new categorical
            variable which is the key in the dictionary
        x_reals: order of continuous variables in tensor passed to forward function
        x_categoricals: order of categorical variables in tensor passed to forward function
        hidden_continuous_size: default for hidden size for processing continous variables (similar to categorical
            embedding size)
        hidden_continuous_sizes: dictionary mapping continuous input indices to sizes for variable selection
            (fallback to hidden_continuous_size if index is not in dictionary)
        embedding_sizes: dictionary mapping (string) indices to tuple of number of categorical classes and
            embedding size
        embedding_paddings: list of indices for embeddings which transform the zero's embedding to a zero vector
        embedding_labels: dictionary mapping (string) indices to list of categorical labels
        learning_rate: learning rate
        log_interval: log predictions every x batches, do not log if 0 or less, log interpretation if > 0. If < 1.0
            , will log multiple entries per batch. Defaults to -1.
        log_val_interval: frequency with which to log validation set metrics, defaults to log_interval
        log_gradient_flow: if to log gradient flow, this takes time and should be only done to diagnose training
            failures
        reduce_on_plateau_patience (int): patience after which learning rate is reduced by a factor of 10
        monotone_constaints (Dict[str, int]): dictionary of monotonicity constraints for continuous decoder
            variables mapping
            position (e.g. ``"0"`` for first position) to constraint (``-1`` for negative and ``+1`` for positive,
            larger numbers add more weight to the constraint vs. the loss but are usually not necessary).
            This constraint significantly slows down training. Defaults to {}.
        share_single_variable_networks (bool): if to share the single variable networks between the encoder and
            decoder. Defaults to False.
        logging_metrics (nn.ModuleList[MultiHorizonMetric]): list of metrics that are logged during training.
            Defaults to nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE()]).
        **kwargs: additional arguments to :py:class:`~BaseModel`.
    """
    if logging_metrics is None:
        logging_metrics = nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE()])
    if loss is None:
        loss = QuantileLoss()
    self.save_hyperparameters()
    # store loss function separately as it is a module
    assert isinstance(loss, MultiHorizonMetric), (
        "Loss has to of class `MultiHorizonMetric`"
    )
    super().__init__(loss=loss, logging_metrics=logging_metrics, **kwargs)

    # processing inputs
    # embeddings
    self.input_embeddings = MultiEmbedding(
        embedding_sizes=self.hparams.embedding_sizes,
        categorical_groups=self.hparams.categorical_groups,
        embedding_paddings=self.hparams.embedding_paddings,
        x_categoricals=self.hparams.x_categoricals,
        max_embedding_size=self.hparams.hidden_size,
    )

    # continuous variable processing
    self.prescalers = nn.ModuleDict(
        {
            name: nn.Linear(
                1,
                self.hparams.hidden_continuous_sizes.get(
                    name, self.hparams.hidden_continuous_size
                ),
            )
            for name in self.hparams.x_reals
        }
    )

    # variable selection
    # variable selection for static variables
    static_input_sizes = {
        name: self.hparams.embedding_sizes[name][1]
        for name in self.hparams.static_categoricals
    }
    static_input_sizes.update(
        {
            name: self.hparams.hidden_continuous_sizes.get(
                name, self.hparams.hidden_continuous_size
            )
            for name in self.hparams.static_reals
        }
    )
    self.static_variable_selection = VariableSelectionNetwork(
        input_sizes=static_input_sizes,
        hidden_size=self.hparams.hidden_size,
        input_embedding_flags={name: True for name in self.hparams.static_categoricals},
        dropout=self.hparams.dropout,
        prescalers=self.prescalers,
    )

    # variable selection for encoder and decoder
    encoder_input_sizes = {
        name: self.hparams.embedding_sizes[name][1]
        for name in self.hparams.time_varying_categoricals_encoder
    }
    encoder_input_sizes.update(
        {
            name: self.hparams.hidden_continuous_sizes.get(
                name, self.hparams.hidden_continuous_size
            )
            for name in self.hparams.time_varying_reals_encoder
        }
    )

    decoder_input_sizes = {
        name: self.hparams.embedding_sizes[name][1]
        for name in self.hparams.time_varying_categoricals_decoder
    }
    decoder_input_sizes.update(
        {
            name: self.hparams.hidden_continuous_sizes.get(
                name, self.hparams.hidden_continuous_size
            )
            for name in self.hparams.time_varying_reals_decoder
        }
    )

    # create single variable grns that are shared across decoder and encoder
    if self.hparams.share_single_variable_networks:
        self.shared_single_variable_grns = nn.ModuleDict()
        for name, input_size in encoder_input_sizes.items():
            self.shared_single_variable_grns[name] = GatedResidualNetwork(
                input_size,
                min(input_size, self.hparams.hidden_size),
                self.hparams.hidden_size,
                self.hparams.dropout,
            )
        for name, input_size in decoder_input_sizes.items():
            if name not in self.shared_single_variable_grns:
                self.shared_single_variable_grns[name] = GatedResidualNetwork(
                    input_size,
                    min(input_size, self.hparams.hidden_size),
                    self.hparams.hidden_size,
                    self.hparams.dropout,
                )

    self.encoder_variable_selection = VariableSelectionNetwork(
        input_sizes=encoder_input_sizes,
        hidden_size=self.hparams.hidden_size,
        input_embedding_flags={
            name: True for name in self.hparams.time_varying_categoricals_encoder
        },
        dropout=self.hparams.dropout,
        context_size=self.hparams.hidden_size,
        prescalers=self.prescalers,
        single_variable_grns={}
        if not self.hparams.share_single_variable_networks
        else self.shared_single_variable_grns,
    )

    self.decoder_variable_selection = VariableSelectionNetwork(
        input_sizes=decoder_input_sizes,
        hidden_size=self.hparams.hidden_size,
        input_embedding_flags={
            name: True for name in self.hparams.time_varying_categoricals_decoder
        },
        dropout=self.hparams.dropout,
        context_size=self.hparams.hidden_size,
        prescalers=self.prescalers,
        single_variable_grns={}
        if not self.hparams.share_single_variable_networks
        else self.shared_single_variable_grns,
    )

    # static encoders
    # for variable selection
    self.static_context_variable_selection = GatedResidualNetwork(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        output_size=self.hparams.hidden_size,
        dropout=self.hparams.dropout,
    )

    # for hidden state of the lstm
    self.static_context_initial_hidden_lstm = GatedResidualNetwork(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        output_size=self.hparams.hidden_size,
        dropout=self.hparams.dropout,
    )

    # for cell state of the lstm
    self.static_context_initial_cell_lstm = GatedResidualNetwork(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        output_size=self.hparams.hidden_size,
        dropout=self.hparams.dropout,
    )

    # for post lstm static enrichment
    self.static_context_enrichment = GatedResidualNetwork(
        self.hparams.hidden_size,
        self.hparams.hidden_size,
        self.hparams.hidden_size,
        self.hparams.dropout,
    )

    # lstm encoder (history) and decoder (future) for local processing
    self.lstm_encoder = nn.LSTM(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        num_layers=self.hparams.lstm_layers,
        dropout=self.hparams.dropout if self.hparams.lstm_layers > 1 else 0,
        batch_first=True,
    )

    self.lstm_decoder = nn.LSTM(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        num_layers=self.hparams.lstm_layers,
        dropout=self.hparams.dropout if self.hparams.lstm_layers > 1 else 0,
        batch_first=True,
    )

    # skip connection for lstm
    self.post_lstm_gate_encoder = GatedLinearUnit(
        self.hparams.hidden_size, dropout=self.hparams.dropout
    )
    self.post_lstm_gate_decoder = self.post_lstm_gate_encoder
    # self.post_lstm_gate_decoder = GatedLinearUnit(self.hparams.hidden_size, dropout=self.hparams.dropout)
    self.post_lstm_add_norm_encoder = AddNorm(
        self.hparams.hidden_size, trainable_add=False
    )
    # self.post_lstm_add_norm_decoder = AddNorm(self.hparams.hidden_size, trainable_add=True)
    self.post_lstm_add_norm_decoder = self.post_lstm_add_norm_encoder

    # static enrichment and processing past LSTM
    self.static_enrichment = GatedResidualNetwork(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        output_size=self.hparams.hidden_size,
        dropout=self.hparams.dropout,
        context_size=self.hparams.hidden_size,
    )

    # attention for long-range processing
    self.multihead_attn = InterpretableMultiHeadAttention(
        d_model=self.hparams.hidden_size,
        n_head=self.hparams.attention_head_size,
        dropout=self.hparams.dropout,
    )
    self.post_attn_gate_norm = GateAddNorm(
        self.hparams.hidden_size, dropout=self.hparams.dropout, trainable_add=False
    )
    self.pos_wise_ff = GatedResidualNetwork(
        self.hparams.hidden_size,
        self.hparams.hidden_size,
        self.hparams.hidden_size,
        dropout=self.hparams.dropout,
    )

    # output processing -> no dropout at this late stage
    self.pre_output_gate_norm = GateAddNorm(
        self.hparams.hidden_size, dropout=None, trainable_add=False
    )

    self.output_layer = nn.Linear(self.hparams.hidden_size, self.hparams.output_size)


# 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 __init__(
    self,
    hidden_size: int = 16,
    lstm_layers: int = 1,
    dropout: float = 0.1,
    output_size: int = 7,
    loss: MultiHorizonMetric = None,
    attention_head_size: int = 4,
    max_encoder_length: int = 10,
    static_categoricals: List[str] = [],
    static_reals: List[str] = [],
    time_varying_categoricals_encoder: List[str] = [],
    time_varying_categoricals_decoder: List[str] = [],
    categorical_groups: Dict[str, List[str]] = {},
    time_varying_reals_encoder: List[str] = [],
    time_varying_reals_decoder: List[str] = [],
    x_reals: List[str] = [],
    x_categoricals: List[str] = [],
    hidden_continuous_size: int = 8,
    hidden_continuous_sizes: Dict[str, int] = {},
    embedding_sizes: Dict[str, Tuple[int, int]] = {},
    embedding_paddings: List[str] = [],
    embedding_labels: Dict[str, np.ndarray] = {},
    learning_rate: float = 1e-3,
    log_interval: Union[int, float] = -1,
    log_val_interval: Union[int, float] = None,
    log_gradient_flow: bool = False,
    reduce_on_plateau_patience: int = 1000,
    monotone_constaints: Dict[str, int] = {},
    share_single_variable_networks: bool = False,
    logging_metrics: nn.ModuleList = None,
    **kwargs,
):
    """
    Temporal Fusion Transformer for forecasting timeseries - use its :py:meth:`~from_dataset` method if possible.

    Implementation of the article
    `Temporal Fusion Transformers for Interpretable Multi-horizon Time Series
    Forecasting <https://arxiv.org/pdf/1912.09363.pdf>`_. The network outperforms DeepAR by Amazon by 36-69%
    in benchmarks.

    Enhancements compared to the original implementation (apart from capabilities added through base model
    such as monotone constraints):

    * static variables can be continuous
    * multiple categorical variables can be summarized with an EmbeddingBag
    * variable encoder and decoder length by sample
    * categorical embeddings are not transformed by variable selection network (because it is a redundant operation)
    * variable dimension in variable selection network are scaled up via linear interpolation to reduce
      number of parameters
    * non-linear variable processing in variable selection network can be shared among decoder and encoder
      (not shared by default)

    Tune its hyperparameters with
    :py:func:`~pytorch_forecasting.models.temporal_fusion_transformer.tuning.optimize_hyperparameters`.

    Args:

        hidden_size: hidden size of network which is its main hyperparameter and can range from 8 to 512
        lstm_layers: number of LSTM layers (2 is mostly optimal)
        dropout: dropout rate
        output_size: number of outputs (e.g. number of quantiles for QuantileLoss)
        loss: loss function taking prediction and targets
        attention_head_size: number of attention heads (4 is a good default)
        max_encoder_length: length to encode (can be far longer than the decoder length but does not have to be)
        static_categoricals: integer of positions of static categorical variables
        static_reals: integer of positions of static continuous variables
        time_varying_categoricals_encoder: integer of positions of categorical variables for encoder
        time_varying_categoricals_decoder: integer of positions of categorical variables for decoder
        time_varying_reals_encoder: integer of positions of continuous variables for encoder
        time_varying_reals_decoder: integer of positions of continuous variables for decoder
        categorical_groups: dictionary where values
            are list of categorical variables that are forming together a new categorical
            variable which is the key in the dictionary
        x_reals: order of continuous variables in tensor passed to forward function
        x_categoricals: order of categorical variables in tensor passed to forward function
        hidden_continuous_size: default for hidden size for processing continous variables (similar to categorical
            embedding size)
        hidden_continuous_sizes: dictionary mapping continuous input indices to sizes for variable selection
            (fallback to hidden_continuous_size if index is not in dictionary)
        embedding_sizes: dictionary mapping (string) indices to tuple of number of categorical classes and
            embedding size
        embedding_paddings: list of indices for embeddings which transform the zero's embedding to a zero vector
        embedding_labels: dictionary mapping (string) indices to list of categorical labels
        learning_rate: learning rate
        log_interval: log predictions every x batches, do not log if 0 or less, log interpretation if > 0. If < 1.0
            , will log multiple entries per batch. Defaults to -1.
        log_val_interval: frequency with which to log validation set metrics, defaults to log_interval
        log_gradient_flow: if to log gradient flow, this takes time and should be only done to diagnose training
            failures
        reduce_on_plateau_patience (int): patience after which learning rate is reduced by a factor of 10
        monotone_constaints (Dict[str, int]): dictionary of monotonicity constraints for continuous decoder
            variables mapping
            position (e.g. ``"0"`` for first position) to constraint (``-1`` for negative and ``+1`` for positive,
            larger numbers add more weight to the constraint vs. the loss but are usually not necessary).
            This constraint significantly slows down training. Defaults to {}.
        share_single_variable_networks (bool): if to share the single variable networks between the encoder and
            decoder. Defaults to False.
        logging_metrics (nn.ModuleList[MultiHorizonMetric]): list of metrics that are logged during training.
            Defaults to nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE()]).
        **kwargs: additional arguments to :py:class:`~BaseModel`.
    """
    if logging_metrics is None:
        logging_metrics = nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE()])
    if loss is None:
        loss = QuantileLoss()
    self.save_hyperparameters()
    # store loss function separately as it is a module
    assert isinstance(loss, MultiHorizonMetric), (
        "Loss has to of class `MultiHorizonMetric`"
    )
    super().__init__(loss=loss, logging_metrics=logging_metrics, **kwargs)

    # processing inputs
    # embeddings
    self.input_embeddings = MultiEmbedding(
        embedding_sizes=self.hparams.embedding_sizes,
        categorical_groups=self.hparams.categorical_groups,
        embedding_paddings=self.hparams.embedding_paddings,
        x_categoricals=self.hparams.x_categoricals,
        max_embedding_size=self.hparams.hidden_size,
    )

    # continuous variable processing
    self.prescalers = nn.ModuleDict(
        {
            name: nn.Linear(
                1,
                self.hparams.hidden_continuous_sizes.get(
                    name, self.hparams.hidden_continuous_size
                ),
            )
            for name in self.reals
        }
    )

    # variable selection
    # variable selection for static variables
    static_input_sizes = {
        name: self.hparams.embedding_sizes[name][1]
        for name in self.hparams.static_categoricals
    }
    static_input_sizes.update(
        {
            name: self.hparams.hidden_continuous_sizes.get(
                name, self.hparams.hidden_continuous_size
            )
            for name in self.hparams.static_reals
        }
    )
    self.static_variable_selection = VariableSelectionNetwork(
        input_sizes=static_input_sizes,
        hidden_size=self.hparams.hidden_size,
        input_embedding_flags={name: True for name in self.hparams.static_categoricals},
        dropout=self.hparams.dropout,
        prescalers=self.prescalers,
    )

    # variable selection for encoder and decoder
    encoder_input_sizes = {
        name: self.hparams.embedding_sizes[name][1]
        for name in self.hparams.time_varying_categoricals_encoder
    }
    encoder_input_sizes.update(
        {
            name: self.hparams.hidden_continuous_sizes.get(
                name, self.hparams.hidden_continuous_size
            )
            for name in self.hparams.time_varying_reals_encoder
        }
    )

    decoder_input_sizes = {
        name: self.hparams.embedding_sizes[name][1]
        for name in self.hparams.time_varying_categoricals_decoder
    }
    decoder_input_sizes.update(
        {
            name: self.hparams.hidden_continuous_sizes.get(
                name, self.hparams.hidden_continuous_size
            )
            for name in self.hparams.time_varying_reals_decoder
        }
    )

    # create single variable grns that are shared across decoder and encoder
    if self.hparams.share_single_variable_networks:
        self.shared_single_variable_grns = nn.ModuleDict()
        for name, input_size in encoder_input_sizes.items():
            self.shared_single_variable_grns[name] = GatedResidualNetwork(
                input_size,
                min(input_size, self.hparams.hidden_size),
                self.hparams.hidden_size,
                self.hparams.dropout,
            )
        for name, input_size in decoder_input_sizes.items():
            if name not in self.shared_single_variable_grns:
                self.shared_single_variable_grns[name] = GatedResidualNetwork(
                    input_size,
                    min(input_size, self.hparams.hidden_size),
                    self.hparams.hidden_size,
                    self.hparams.dropout,
                )

    self.encoder_variable_selection = VariableSelectionNetwork(
        input_sizes=encoder_input_sizes,
        hidden_size=self.hparams.hidden_size,
        input_embedding_flags={
            name: True for name in self.hparams.time_varying_categoricals_encoder
        },
        dropout=self.hparams.dropout,
        context_size=self.hparams.hidden_size,
        prescalers=self.prescalers,
        single_variable_grns={}
        if not self.hparams.share_single_variable_networks
        else self.shared_single_variable_grns,
    )

    self.decoder_variable_selection = VariableSelectionNetwork(
        input_sizes=decoder_input_sizes,
        hidden_size=self.hparams.hidden_size,
        input_embedding_flags={
            name: True for name in self.hparams.time_varying_categoricals_decoder
        },
        dropout=self.hparams.dropout,
        context_size=self.hparams.hidden_size,
        prescalers=self.prescalers,
        single_variable_grns={}
        if not self.hparams.share_single_variable_networks
        else self.shared_single_variable_grns,
    )

    # static encoders
    # for variable selection
    self.static_context_variable_selection = GatedResidualNetwork(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        output_size=self.hparams.hidden_size,
        dropout=self.hparams.dropout,
    )

    # for hidden state of the lstm
    self.static_context_initial_hidden_lstm = GatedResidualNetwork(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        output_size=self.hparams.hidden_size,
        dropout=self.hparams.dropout,
    )

    # for cell state of the lstm
    self.static_context_initial_cell_lstm = GatedResidualNetwork(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        output_size=self.hparams.hidden_size,
        dropout=self.hparams.dropout,
    )

    # for post lstm static enrichment
    self.static_context_enrichment = GatedResidualNetwork(
        self.hparams.hidden_size,
        self.hparams.hidden_size,
        self.hparams.hidden_size,
        self.hparams.dropout,
    )

    # lstm encoder (history) and decoder (future) for local processing
    self.lstm_encoder = nn.LSTM(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        num_layers=self.hparams.lstm_layers,
        dropout=self.hparams.dropout if self.hparams.lstm_layers > 1 else 0,
        batch_first=True,
    )

    self.lstm_decoder = nn.LSTM(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        num_layers=self.hparams.lstm_layers,
        dropout=self.hparams.dropout if self.hparams.lstm_layers > 1 else 0,
        batch_first=True,
    )

    # skip connection for lstm
    self.post_lstm_gate_encoder = GatedLinearUnit(
        self.hparams.hidden_size, dropout=self.hparams.dropout
    )
    self.post_lstm_gate_decoder = self.post_lstm_gate_encoder
    # self.post_lstm_gate_decoder = GatedLinearUnit(self.hparams.hidden_size, dropout=self.hparams.dropout)
    self.post_lstm_add_norm_encoder = AddNorm(
        self.hparams.hidden_size, trainable_add=False
    )
    # self.post_lstm_add_norm_decoder = AddNorm(self.hparams.hidden_size, trainable_add=True)
    self.post_lstm_add_norm_decoder = self.post_lstm_add_norm_encoder

    # static enrichment and processing past LSTM
    self.static_enrichment = GatedResidualNetwork(
        input_size=self.hparams.hidden_size,
        hidden_size=self.hparams.hidden_size,
        output_size=self.hparams.hidden_size,
        dropout=self.hparams.dropout,
        context_size=self.hparams.hidden_size,
    )

    # attention for long-range processing
    self.multihead_attn = InterpretableMultiHeadAttention(
        d_model=self.hparams.hidden_size,
        n_head=self.hparams.attention_head_size,
        dropout=self.hparams.dropout,
    )
    self.post_attn_gate_norm = GateAddNorm(
        self.hparams.hidden_size, dropout=self.hparams.dropout, trainable_add=False
    )
    self.pos_wise_ff = GatedResidualNetwork(
        self.hparams.hidden_size,
        self.hparams.hidden_size,
        self.hparams.hidden_size,
        dropout=self.hparams.dropout,
    )

    # output processing -> no dropout at this late stage
    self.pre_output_gate_norm = GateAddNorm(
        self.hparams.hidden_size, dropout=None, trainable_add=False
    )

    self.output_layer = nn.Linear(self.hparams.hidden_size, self.hparams.output_size)