How to properly move submodules to GPU?
I've coded up a TransformerEncoder that relies on submodules. Specifically, I have a main lightning module (MainTransfomer.py) which has 2 sub (regular torch)modules. 1 is BertModel and 1 is a custom TransformerEncoderLayer.
The TransformerEncoderLayer has 2 submodules of its own named MultiHeadAttn and PosFeedForward. When I try to run my code via the lightning trainer, it seems that all my tensors/params are being moved to GPU except for when it tries to do operations within the MultiHeadAttn and PosFeedForward modules. Specifically it seems the weight matrix of a nn.Linear() within MultiHeadAttn is still sitting on cpu.
So my question boils down to what is the correct way to ensure all submodules are moved to the (correct) GPU within the lightning framework? In normal PyTorch I would just explicitly call .to(device) but I've read that it is not recommended to do this within lightning.
if I explicitly set .to(device) on my TransformerEncoderLayer in my MainTransformerEncoder init I don't run into this issue.
I can supply code if need be, but the general setup is as described above. Main lightning module inits a torch module, which itself has 2 sub torch modules (attn & ffnn). The attn and ffnn modules don't seem to be moved to GPU by the lightning trainer.
Env:

Ubuntu16.04
conda/py3.8
Lightning v0.7.1