Running demo.py with different checkpoints

[emrecncelik]

Hello,

How can I use baseplus and large checkpoints with the demo code? When I provide the model path as either baseplus or large I get the following error (when set to strict=True in load_state_dict, otherwise does not give any errors, prints out modules and stops),

RuntimeError: Error(s) in loading state_dict for OptimizedModule:
        Missing key(s) in state_dict: "_orig_mod.encoder.frontend_a.trunk.conv1.weight", "_orig_mod.encoder.frontend_a.trunk.bn1.weight", 
"_orig_mod.encoder.frontend_a.trunk.bn1.bias", "_orig_mod.encoder.frontend_a.trunk.bn1.running_mean", 
"_orig_mod.encoder.frontend_a.trunk.bn1.running_var", "_orig_mod.encoder.frontend_a.trunk.layer1.0.conv1.weight", 
"_orig_mod.encoder.frontend_a.trunk.layer1.0.bn1.weight", "_orig_mod.encoder.frontend_a.trunk.layer1.0.bn1.bias", 
"_orig_mod.encoder.frontend_a.trunk.layer1.0.bn1.running_mean", "_orig_mod.encoder.frontend_a.trunk.layer1.0.bn1.running_var", 
"_orig_mod.encoder.frontend_a.trunk.layer1.0.conv2.weight", "_orig_mod.encoder.frontend_a.trunk.layer1.0.bn2.weight", 

It's actually longer and prints out all modules with the prefix "_orig_mod". So I am currently using large and baseplus models with the following modification to the demo code,

from hydra.utils import instantiate

ckpt = torch.load(cfg.model.pretrained_model_path, map_location="cpu")
usr = torch.compile(instantiate(cfg.model.obj, cfg))
usr.load_state_dict(ckpt)
usr_mod = usr._orig_mod if hasattr(usr, "_orig_mod") else usr
model = usr_mod.model.backbone

and overriding backbone in the Hydra config by,

python demo.py model/backbone=resnet_transformer_baseplus model.pretrained_model_path=/path/to/baseplus

But I am not entirely sure if this is the correct way. So my question is how can I correctly use baseplus and large models in the demo.py?

Thank you.