Load checkpoint with dynamically created model


In the Lightingmodule docs, the setup hook is described as a possibility to dynamically build a model (instead of initiating in __init__). See the example here.

However, when I load a Lightingmodule checkpoint (load_from_checkpoint) with such a dynamically created model, I get a runtime Error (RuntimeError: Error(s) in loading state_dict …) as setup is not called within the loading procedure.

Is there a way to have dynamically generated models in a Lightningmodule (created in setup) that can be loaded from a checkpoint?