Load pretrained pytorch lightning model

[Djennifer]

I have trained a pytorch lightning model and am now trying to load it in MONAILabel. However, I am stumbling upon some problems. Pytorch lightning saves its weights differently from conventional pytorch. This is why I am in need of more flexibility when it comes to loading my model before running inference. Below you can see how I do it now but eventually, I run into an AttributeError: '_IncompatibleKeys' object has no attribute 'to' from the function _get_network in basic_infer. While if I load the model like this out of MONAILabel it works

    # Network
    self.path = os.path.join(self.model_dir, f"pretrained_{name}.ckpt")

    checkpoint = torch.load(self.path)
    hyper_parameters = checkpoint["hyper_parameters"]['config']['model_params']

    Net = UNet(**hyper_parameters)

    model_weights = checkpoint["state_dict"]

    # # update keys by dropping `_model.`
    for key in list(model_weights):
        model_weights[key.replace("_model.", "")] = model_weights.pop(key)

    self.network = Net.load_state_dict(model_weights)