I think the memory leak of this code comes from that you are constantly appending gradient tensors to the embedding_gradient list.
Since you do not free it after calculation, it remains in the GPU memory.
However, I’m not quite sure about what you try to achieve here.
First, calling gradient.backward() before optimizer.step() is treating the gradient as your loss function.
So you are actually training the parameters to minimize the gradient, is this what you want?
Second, I think that in most NLP tasks it’s common to accumulate the gradient of multiple occurrences for a single token, since a token need information from multiple positions to update its representation.
Yet if you think you really need to retain gradients from specific positions, you should block the gradients by customize a torch.autograd.Function instead of using a hook.