import torch
import torch.nn as nn

class CustomEmbeddingBag(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, mode, include_last_offset):
        super(CustomEmbeddingBag, self).__init__()
        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
        self.mode = mode
        self.include_last_offset = include_last_offset

    def forward(self, indices, offsets):
        if self.include_last_offset:
            # Correctly compute the last bag's end
            last_bag_end = offsets[-1]
            indices = indices[:last_bag_end]
        embeddings = self.embedding(indices)
        if self.mode == "sum":
            output = torch.zeros((len(offsets) - 1, embeddings.shape[1]), device=embeddings.device)
            start =            for i in range(len(offsets) - 1):
                end = offsets[i + 1]
                output[i] = torch.sum(embeddings[start:end], dim=0)
                start = end
        return output

# Usage
weight = torch.tensor(
    [[1.,  2.],
     [3.,  4.],
     [5.,  6.],
     [7.,  8.],
     [99., 99.]],  # row 4 — should NOT appear in any bag
    dtype=torch.float32,
)
indices = torch.tensor([0, 1, 2, 3, 4], dtype=torch.long)
offsets = torch.tensor([0, 2, 4],       dtype=torch.long)

eb_cpu = CustomEmbeddingBag(5, 2, "sum", True)
eb_cpu.embedding.weight = torch.nn.Parameter(weight.clone())
out_cpu = eb_cpu(indices, offsets)

eb_cuda = CustomEmbeddingBag(5, 2, "sum", True).cuda()
eb_cuda.embedding.weight = torch.nn.Parameter(weight.clone().cuda())
out_cuda = eb_cuda(indices.cuda(), offsets.cuda())

print(f"CPU  bag1: {out_cpu[1].tolist()}")   # [