I’m currently trying to find a faster way to process many (~100) categorical embeddings in the forward pass.
As the features have distinct vocabulary sizes and embedding dimensions, each requires its own embedding layer. Currently, I initialize them as nn.ModuleDict (emb_model in the snippet below), where the keys correspond to the names of the categorical features and the values correspond to the embedding layers (each with a weight tensor of shape (size_vocab, d_emb)).
In the forward pass, I loop over the ~100 .items() of the ModuleDict and provide the corresponding features as inputs (X is a dictionary with the same keys as the ModuleDict) to each layer:
def forward(self, X):

    concat_list = []
    for key, layer in self.emb_model.items():
        x = layer(X[key])
        concat_list.append(x)

    x = torch.cat(concat_list, dim=1)
...

However, this solution is about 3x slower than a equivalent model using the Keras/TF functional API, which I presume is due to the loop being executed in every forward pass.
Is there any better way to handle large numbers of categorical features in PyTorch?