I’m trying to hack the nn.EmbeddingBag module to get the desired effect:
import torch
import torch.nn as nn
from torch.autograd import Variable


class TensorReducer(nn.Module):
    
    def __init__(self):
        super(TensorReducer, self).__init__()
        self.emb = nn.EmbeddingBag(0,0)
        
    def set_embs(self, embeddings):
        self.emb.weight = nn.Parameter(embeddings.data)
        
    def get_dummy_input(self, size):
        return Variable(torch.LongTensor(list(range(size))))
    
    def forward(self, inputs, offsets):
        self.set_embs(inputs)
        dummy = self.get_dummy_input(inputs.size(0))
        return self.emb(dummy, offsets)
    

red = TensorReducer()

inputs = Variable(torch.FloatTensor([[1,1,1],[2,2,2],[3,3,3],
                                     [4,4,4],[5,5,5],[6,6,6]]))
offsets = Variable(torch.LongTensor([0,4]))
red(inputs, offsets)
#  Variable containing:
#   2.5000  2.5000  2.5000
#   5.5000  5.5000  5.5000
#[torch.FloatTensor of size 2x3]

The issue is that when I’m setting the internal embedding matrix inside the EmbeddingBag I’m breaking the ability of the gradient to backpropagate in earlier module. Any solution for this?