embedding manual control location CPU vs GPU

[johngrabner]

I would like to create an embedding that does not fit in the GPU memory
but can fit in the CPU memory.

Select the subset for a batch, send it to the GPU at the start of mini-batch.

GPU_tensor = embedding(idx)

Then at the end of training update the CPU embedding from the GPU embedding.

I am using
pl.Trainer( gpus=[0,1], distributed_backend='ddp')

and probably will need accumulate_grad_batches

Any idea for how to do this ?

[carmocca]

You should be able to move any tensors as you see fit with tensor.to(...)