Yes currently I am using a weight mask to invalidate the data of the specific indices per batch. Kind of like this:
    def invalidate_embeddings(self, max_ent_id):
        """
        Invalidate the entity embeddings
        :param max_ent_id: int
        :return:
        """
        vocab_size = self.embedding.weight.data.size(0)
        ids = torch.arange(0, vocab_size)
        mask = ids > max_ent_id
        mask = mask.unsqueeze(1)
        weight = self.embedding.weight.data
        weight = weight * mask
        self.embedding.weight.data = weight

Till now it’s good because I kept a block of indices at the beginning of the embedding matrix. It gets a bit trickier if the indices are not contiguous: then we need loops and the training will be slower.