Lists of PyTorch Lightning sub-models don't get transferred to GPU
When using PyTorch Lightning on CPU, everything works fine. However when using GPUs, I get a RuntimeError: Expected all tensors to be on the same device.
It seems that the trouble comes from the model using a list of sub-models which don't get passed to the GPU:
class LambdaLayer(LightningModule):
    def __init__(self, fun):
        super(LambdaLayer, self).__init__()
        self.fun = fun

    def forward(self, x):
        return self.fun(x)

class TorchModel(LightningModule):
    def __init__(self):
        super(TorchModel, self).__init__()
        self.cat_layers = [TorchCatEmbedding(cat) for cat in columns_to_embed]
        self.num_layers = [LambdaLayer(lambda x: x[:, idx:idx+1]) for _, idx in numeric_columns]
        self.ffo = TorchFFO(len(self.num_layers) + sum([embed_dim(l) for l in self.cat_layers]), y.shape[1])
        self.softmax = torch.nn.Softmax(dim=1)

model = TorchModel()
trainer = Trainer(gpus=-1)

Before running trainer(model):
&gt;&gt;&gt; model.device
device(type='cpu')

&gt;&gt;&gt; model.ffo.device
device(type='cpu')

&gt;&gt;&gt; model.cat_layers[0].device
device(type='cpu')

After running trainer(model):
&gt;&gt;&gt; model.device
device(type='cuda', index=0) # &lt;---- correct

&gt;&gt;&gt; model.ffo.device
device(type='cuda', index=0) # &lt;---- correct

&gt;&gt;&gt; model.cat_layers[0].device
device(type='cpu') # &lt;---- still showing 'cpu'

Apparently, PyTorch Lightning is not able to transfer the lists of sub-models to the GPU. How to proceed so that the entire model, including list of sub-models (cat_layers and num_layers) is transferred to the GPU?

Submodules contained in lists are not registered and can't be transformed as is.
You need to use ModuleList instead, i.e.:
...
from torch.nn import ModuleList
...

class TorchModel(LightningModule):
    def __init__(self):
        super(TorchModel, self).__init__()
        self.cat_layers = ModuleList([TorchCatEmbedding(cat) for cat in columns_to_embed])
        self.num_layers = ModuleList([LambdaLayer(lambda x: x[:, idx:idx+1]) for _, idx in numeric_columns])
        self.ffo = TorchFFO(len(self.num_layers) + sum([embed_dim(l) for l in self.cat_layers]), y.shape[1])
        self.softmax = torch.nn.Softmax(dim=1)

edit: I'm not sure what the Lightning equivalent is, or if one such exists, see also PyTorch Lightning - LightningModule for ModuleList / ModuleDict?
