Cannot load saved MTLSD model

[Mobinapournemat]

Hello,

After training the MTLSD model using the provided notebook, I saved the checkpoints. I'm trying to load the model using saved checkpoints to predict affs and lsds using this code (from segment.ipynb notebook) :

raw, pred_lsds, pred_affs = predict(checkpoint, raw_file, raw_dataset)

but I get this error:

Traceback (most recent call last): File "/usr/local/lib/python3.10/dist-packages/gunpowder/nodes/batch_provider.py", line 193, in request_batch batch = self.provide(upstream_request) File "/usr/local/lib/python3.10/dist-packages/gunpowder/nodes/batch_filter.py", line 148, in provide dependencies = self.prepare(request) File "/usr/local/lib/python3.10/dist-packages/gunpowder/nodes/generic_predict.py", line 116, in prepare self.start() File "/usr/local/lib/python3.10/dist-packages/gunpowder/torch/nodes/predict.py", line 107, in start self.model.load_state_dict(checkpoint["model_state_dict"]) File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 2152, in load_state_dict raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format( RuntimeError: Error(s) in loading state_dict for MtlsdModel: Missing key(s) in state_dict: "unet.r_up.0.0.up.weight", "unet.r_up.0.0.up.bias", "unet.r_up.0.1.up.weight", "unet.r_up.0.1.up.bias".

I also tried training the model with the default dataset used in the notebooks but got the same error again. I appreciate any help or suggestions.