Training composite models on the GPU

[rdisipio]

Hi,

I'm having issues training my model due to some CPU/GPU mismatch. Before I call this a bug, I though the Trainer class handled the upload to the GPU memory by itself. However, I suspect this works only if the model is self-contained. Can anyone confirm one can do something like this?

class ModelA(pl.LightningModule):
  def __init__(self):
       super(ModelA, self).__init__()
  blablabla

class ModelB(pl.LightningModule):
  def __init__(self):
       super(ModelB, self).__init__()
       self.layer = ModelA()

  def forward(self, x):
     self.layer(x)
  blablabla

[rdisipio]

I think I've found a solution. I had a list of sublayers in a plain python list. Encapsulating the list in a nn.ModuleList object solved the issue.