Thanks for the advice @ptrblck. I have a couple of questions in that case,

What’s the running time of updating a batch of embedding vectors within my embedding matrix at the end of each training batch? Is it O(batch_size) or O(matrix_size)?
Would it be faster to load a “chunk” (maybe 100 batches) of embeddings into gpu in each Dataloader iteration and do a for-loop over batches in a chunk for training? (Assuming you’re ok with some embeddings being stale in the chunk)