Is it ok to define model classes outside of the class that inherits from LightningModule?

[EvanZ]

I'm getting the following runtime error:

Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

I have removed all code that explicitly moves tensors to the gpu or cpu so I'm not sure what is going on. But one question I have that I thought might be related is whether I can define a class that is external to the class that inherits from Lightning Module, such as this:

class ExternalModel(torch.nn.module):
    def __init__(self):
        does stuff
    forward():
        does stuff

class MyLightningModule(pl.LightningModule):
    def __init__(self):
        model = ExternalModel()

   def forward(self):
       does stuff

I've been reading through docs trying to find out if this is allowed. I could imagine if one has a lot of existing code, it could get awkward quite fast to have to contain all of it within the LightningModule class.

[tchaton]

Hey @EvanZ,

Did you check ExternalModel is a submodule of the LightningModule.
From your pseudo code. it is missing setting model to self.

class MyLightningModule(pl.LightningModule):
    def __init__(self):
        self.model = ExternalModel()

   def forward(self):
       self.model(...)

[EvanZ]

Yes, @tchaton I am setting it to self.model in my code.