Training sending undesired tensor to GPU

[felipemello1]

🐛 Bug

Embedding vector, defined outside of training loop, is sent to GPU if trainer has gpu > 0 when trainer starts to run.

Please reproduce using the BoringModel

To Reproduce

https://colab.research.google.com/gist/fmellomascarenhas/3054651c4081f1489ff8d56fe0636a1d/the-boringmodel.ipynb

Expected behavior

My embedding is too big to fit inside of my GPU. I want to be able to call the Embedding in the CPU, and then send the selected embeddings to GPU for training.

[carmocca]

Lightning automatically moves the model to the appropriate device before training starts. Since the embeddings are considered part of the model, they will get moved too.

You can move it back to cpu by overriding the on_fit_start hook (or any other after the setup hook).

        def on_fit_start(self, trainer, pl_module):
            self.emb = self.emb.cpu()

You can see the hooks and their order here

[carmocca]

Duplicate of #3698

Closing, feel free to open if you have any other issues