Implementing LightFM model on PyTorch

I’m trying to implement the LightFM model (https://github.com/lyst/lightfm) on PyTorch, and I’m starting with a case where I have users, items and item features. The item embedding is a sum of embeddings of its features. However, I’m facing a computational time problem because each item has a different number of features, and I can’t store them in one tensor.

Here the code.

from typing import List

import numpy as np
import torch


class Model(torch.nn.Module):
    
    def __init__(
            self, 
            embedding_dim: int,
            num_users: int, 
            num_item_features: int,
            item_features: List[np.array],
            device: str="cpu",
    ):
        super().__init__()
        self.user_embeddings = torch.nn.Embedding(
            num_embeddings=num_users,
            embedding_dim=embedding_dim,
            sparse=True,
            device=device,
        )
        self.item_feature_embeddings = torch.nn.Embedding(
            num_embeddings=num_item_features,
            embedding_dim=embedding_dim,
            sparse=True,
            device=device,
        )
        self.item_features = [
            torch.LongTensor(features).to(device)
            for features in item_features
        ]
        torch.nn.init.xavier_uniform_(self.user_embeddings.weight)
        torch.nn.init.xavier_uniform_(self.item_feature_embeddings.weight)
        self.device = device
        
    def item_embeddings(self, item_ids: np.array) -> torch.FloatTensor:
        item_features = [
            self.item_features[item_id]
            for item_id in item_ids
        ]
        item_embeddings = [
            self.item_feature_embeddings(features).sum(axis=0) 
            for features in item_features
        ]
        return torch.stack(item_embeddings)
        
    def forward(self, user_ids: np.array, item_ids: np.array) -> torch.FloatTensor:
        user_ids = torch.LongTensor(user_ids).to(self.device)
        user_embeddings = self.user_embeddings(user_ids)
        item_embeddings = self.item_embeddings(item_ids)
        return torch.mm(user_embeddings, item_embeddings.t())


NUM_USERS = 1_000_000
NUM_ITEMS = 1_000_000
NUM_ITEM_FEATURES = 2_000_000
EMBEDDING_DIM = 128
DEVICE = 'cuda'

item_features = [
    np.random.randint(
        low=0,
        high=num_item_features, 
        size=np.random.binomial(n=1000, p=0.01) + 1,
    )
    for _ in range(NUM_ITEMS)
]
model = Model(
    embedding_dim=EMBEDDING_DIM,
    num_users=NUM_USERS,
    num_item_features=NUM_ITEM_FEATURES,
    item_features=item_features,
    device=DEVICE,
)
user_ids = np.random.randint(low=0, high=NUM_USERS, size=10000)
item_ids = np.random.randint(low=0, high=NUM_ITEMS, size=10000)

scores = model(user_ids=user_ids, item_ids=item_ids)

But it take a lot of time. The line

        item_embeddings = [
            self.item_feature_embeddings(features).sum(axis=0) 
            for features in item_features
        ]

take more than 95% of time.
How can I store item features and use them so that it will be more efficient?