Thank you for spending to fix bug. I found that memory usage becomes increase when updating the embedding. The problem was in model code
    self.embedding = nn.Embedding.from_pretrained(embedding)
    if embedding is not None:
        self.embedding.weight = nn.Parameter(embedding)
    self.embedding.weight.requires_grad = False  #True

Memory problem was triggered when requires_grad =  True when using the pretrained embedding.