def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        """
        Parameters
        ----------
        x
            input tensor of shape batch x (optional) time x categoricals in the order of ``variable_names``.

        Returns
        -------
        dict
            dictionary of category names to embeddings of shape batch x (optional) time x embedding_size if
            ``embedding_size`` is given as dictionary.
        """
        return {
            name: self.embeddings[name](x[..., i])
            for i, name in enumerate(self.variable_names)
        }

        Parameters
        ----------
        x
            input tensor of shape batch x (optional) time x categoricals in the order of ``variable_names``.

        Returns
        -------
        dict
            dictionary of category names to embeddings of shape batch x (optional) time x embedding_size if
            ``embedding_size`` is given as dictionary.
        
darts/models/forecasting/tft_submodels.py