I’m not seeing the advice of leaving the embedding on the CPU in the linked issue. However, if you would like to do it, you could just call model.embedding.cpu(), where .embedding refers to the attribute name of your nn.Embedding layer. In the forward method you would then have to pass a CPU tensor to the embedding, push the output to the GPU and pass it to the next layer (which should have its parameters on the GPU).


Yes, that should work. Could you post your model definition and a code snippet which reproduces this error, so that we can have a look?