How to put nn.Embedding on CPU

1

I am working with a seq2seq model. Since the vocabulary size is large. I want to put the nn.Embedding layer on the main memory, not on GPU. Is there anyway to achieve that? Thanks.

2

yes, just call .cpu() on the layer, so that it is on the CPU side of things.
You can call .cuda() on the output of the layer and then send it to your LSTM that resides on the GPU.

3

@cswhjiang
Did you manage to implement it. I am confused because when i put model.cuda(), won’t the embeddings also get to gpu as well.
If you did manage to implement it, would you mind sharing a snippet of it.

Thanks