Can't Load Nested Models from Checkpoints using load_from_checkpoint()

[malfonsoarquimea]

I'm dealing with a situation where I have a primary model that contains another model nested within it. The structure of my code is as follows:

def PrimaryModel(..):
   def init(...)
      nested_model=NestedModel(nested_model_ckpt_path,...)
      self.mlp1 = nested_model.mlp1
      self.mlp2 = nested_model.mlp2

When I attempt to use PrimaryModel.load_from_checkpoint(primary_model_ckpt_path), an error occurs during the initialization of NestedModel because it can't find its nested_model_ckpt_path. I'm not passing this path because I'm loading PrimaryModel from its own checkpoint.

Additionally, I've confirmed that the state dictionary from the PrimaryModel checkpoint contains the information for mlp1 and mlp2:

state_dict=torch.load(cfg.primary_model_ckpt_path)["state_dict"]

I'm looking for a clean way to load the nested NestedModel from the PrimaryModel checkpoint, since it contains all the necessary information for both models. Any guidance on how I should modify my code to handle this properly would be greatly appreciated.