What about using directly nn.functionnal.embedding_bag:
weights = 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]))
input = Variable(torch.arange(weights.size(0)).long())

nn.functional.embedding_bag(weights, input, offsets, None, 2, False, 'mean')

Note that you need at least pytorch v3 to use it.