Embedding weight is not being updated on training

Below is my model code and expecting to have embedding weight updated as training is progressing but it is not being updated. Only embedding which is being updated is of index = 0 which is ‘unk’ word. Am I missing anything in the model ? Could it be possible someone to help on this ? I am struggling for last few days.

class FeedForwardModel(nn.Module):
    def __init__(self, embedding_matrix, embed_dim):
        super().__init__()
        self.embedding = nn.EmbeddingBag.from_pretrained(torch.FloatTensor(embedding_matrix))
        self.fc1 = nn.Linear(embed_dim, 10)
        self.fc2 = nn.Linear(10, 1)
        self.output = nn.Sigmoid()
        self.init_weights()
        
    def __init__(self, num_embeddings, embed_dim):
        super().__init__()
        self.embedding = nn.EmbeddingBag(num_embeddings, embedding_dim)
        self.fc1 = nn.Linear(embed_dim, 10)
        self.fc2 = nn.Linear(10, 1)
        self.output = nn.Sigmoid()
        self.embedding.weight.data.uniform_(-0.5, 0.5)
        self.embedding.weight.requires_grad = True
        self.init_weights()
                  
        
    def forward(self, input, offsets):
        #print(input)
        #print(self.embedding.weight[1:3])
        embedded = self.embedding(input, offsets)
        h1 = F.tanh(self.fc1(embedded))
        h2 = self.fc2(h1)
        return self.output(h2)
    
    def init_weights(self):
        initrange = 0.5
        self.fc1.weight.data.uniform_(-initrange, initrange)
        self.fc2.weight.data.uniform_(-initrange, initrange)
        self.fc1.bias.data.zero_()
        self.fc2.bias.data.zero_()

Could you check your input for all unique values?
If you are passing all valid embedding indices, you’ll get valid gradients in the complete weight parameter of self.embedding:


num_embeddings = 10
embedding_dim = 100
model = FeedForwardModel(num_embeddings, embedding_dim)

x = torch.arange(num_embeddings)
offsets = torch.tensor([0, 5])

output = model(x, offsets)
output.mean().backward()

print(model.embedding.weight.grad)

However, if you pass e.g. only zeros as input, only weight.grad[0] will be filled with non-zero values.