May I ask how to fix this error? (When using my own trained model to infer)

[silence-tang]

Traceback (most recent call last): File "infer.py", line 144, in user = User(ARCH, DATA, FLAGS.dataset, FLAGS.log, FLAGS.model,FLAGS.split,FLAGS.uncertainty,FLAGS.monte_carlo) File "../../tasks/semantic/modules/user.py", line 70, in init self.model.load_state_dict(w_dict['state_dict'], strict=True) File "/root/miniconda3/envs/salsanext/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1605, in load_state_dict self.class.name, "\n\t".join(error_msgs))) RuntimeError: Error(s) in loading state_dict for DataParallel: Missing key(s) in state_dict: "module.downCntx.conv1.weight", "module.downCntx.conv1.bias", "module.downCntx.conv2.weight", "module.downCntx.conv2.bias", "module.downCntx.bn1.weight", "module.downCntx.bn1.bias", "module.downCntx.bn1.running_mean", "module.downCntx.bn1.running_var", "module.downCntx.conv3.weight", "module.downCntx.conv3.bias", "module.downCntx.bn2.weight", "module.downCntx.bn2.bias", "module.downCntx.bn2.running_mean", "module.downCntx.bn2.running_var", "module.downCntx2.conv1.weight", "module.downCntx2.conv1.bias", "module.downCntx2.conv2.weight", "module.downCntx2.conv2.bias", "module.downCntx2.bn1.weight", "module.downCntx2.bn1.bias", "module.downCntx2.bn1.running_mean", "module.downCntx2.bn1.running_var", "module.downCntx2.conv3.weight", "module.downCntx2.conv3.bias", "module.downCntx2.bn2.weight", "module.downCntx2.bn2.bias", "module.downCntx2.bn2.running_mean", "module.downCntx2.bn2.running_var", "module.downCntx3.conv1.weight", "module.downCntx3.conv1.bias", "module.downCntx3.conv2.weight", "module.downCntx3.conv2.bias", "module.downCntx3.bn1.weight", "module.downCntx3.bn1.bias", "module.downCntx3.bn1.running_mean", "module.downCntx3.bn1.running_var", "module.downCntx3.conv3.weight", "module.downCntx3.conv3.bias", "module.downCntx3.bn2.weight", "module.downCntx3.bn2.bias", "module.downCntx3.bn2.running_mean", "module.downCntx3.bn2.running_var", "module.resBlock1.conv1.weight", "module.resBlock1.conv1.bias", "module.resBlock1.conv2.weight", "module.resBlock1.conv2.bias", "module.resBlock1.bn1.weight", "module.resBlock1.bn1.bias", "module.resBlock1.bn1.running_mean", "module.resBlock1.bn1.running_var", "module.resBlock1.conv3.weight", "module.resBlock1.conv3.bias", "module.resBlock1.bn2.weight", "module.resBlock1.bn2.bias", "module.resBlock1.bn2.running_mean", "module.resBlock1.bn2.running_var", "module.resBlock1.conv4.weight", "module.resBlock1.conv4.bias", "module.resBlock1.bn3.weight", "module.resBlock1.bn3.bias", "module.resBlock1.bn3.running_mean", "module.resBlock1.bn3.running_var", "module.resBlock1.conv5.weight", "module.resBlock1.conv5.bias", "module.resBlock1.bn4.weight", "module.resBlock1.bn4.bias", "module.resBlock1.bn4.running_mean", "module.resBlock1.bn4.running_var", "module.resBlock2.conv1.weight", "module.resBlock2.conv1.bias", "module.resBlock2.conv2.weight", "module.resBlock2.conv2.bias", "module.resBlock2.bn1.weight", "module.resBlock2.bn1.bias", "module.resBlock2.bn1.running_mean", "module.resBlock2.bn1.running_var", "module.resBlock2.conv3.weight", "module.resBlock2.conv3.bias", "module.resBlock2.bn2.weight", "module.resBlock2.bn2.bias", "module.resBlock2.bn2.running_mean", "module.resBlock2.bn2.running_var", "module.resBlock2.conv4.weight", "module.resBlock2.conv4.bias", "module.resBlock2.bn3.weight", "module.resBlock2.bn3.bias", "module.resBlock2.bn3.running_mean", "module.resBlock2.bn3.running_var", "module.resBlock2.conv5.weight", "module.resBlock2.conv5.bias", "module.resBlock2.bn4.weight", "module.resBlock2.bn4.bias", "module.resBlock2.bn4.running_mean", "module.resBlock2.bn4.running_var", "module.resBlock3.conv1.weight", "module.resBlock3.conv1.bias", "module.resBlock3.conv2.weight", "module.resBlock3.conv2.bias", "module.resBlock3.bn1.weight", "module.resBlock3.bn1.bias", "module.resBlock3.bn1.running_mean", "module.resBlock3.bn1.running_var", "module.resBlock3.conv3.weight", "module.resBlock3.conv3.bias", "module.resBlock3.bn2.weight", "module.resBlock3.bn2.bias", "module.resBlock3.bn2.running_mean", "module.resBlock3.bn2.running_var", "module.resBlock3.conv4.weight", "module.resBlock3.conv4.bias", "module.resBlock3.bn3.weight", "module.resBlock3.bn3.bias", "module.resBlock3.bn3.running_mean", "module.resBlock3.bn3.running_var", "module.resBlock3.conv5.weight", "module.resBlock3.conv5.bias", "module.resBlock3.bn4.weight", "module.resBlock3.bn4.bias", "module.resBlock3.bn4.running_mean", "module.resBlock3.bn4.running_var", "module.resBlock4.conv1.weight", "module.resBlock4.conv1.bias", "module.resBlock4.conv2.weight", "module.resBlock4.conv2.bias", "module.resBlock4.bn1.weight", "module.resBlock4.bn1.bias", "module.resBlock4.bn1.running_mean", "module.resBlock4.bn1.running_var", "module.resBlock4.conv3.weight", "module.resBlock4.conv3.bias", "module.resBlock4.bn2.weight", "module.resBlock4.bn2.bias", "module.resBlock4.bn2.running_mean", "module.resBlock4.bn2.running_var", "module.resBlock4.conv4.weight", "module.resBlock4.conv4.bias", "module.resBlock4.bn3.weight", "module.resBlock4.bn3.bias", "module.resBlock4.bn3.running_mean", "module.resBlock4.bn3.running_var", "module.resBlock4.conv5.weight", "module.resBlock4.conv5.bias", "module.resBlock4.bn4.weight", "module.resBlock4.bn4.bias", "module.resBlock4.bn4.running_mean", "module.resBlock4.bn4.running_var", "module.resBlock5.conv1.weight", "module.resBlock5.conv1.bias", "module.resBlock5.conv2.weight", "module.resBlock5.conv2.bias", "module.resBlock5.bn1.weight", "module.resBlock5.bn1.bias", "module.resBlock5.bn1.running_mean", "module.resBlock5.bn1.running_var", "module.resBlock5.conv3.weight", "module.resBlock5.conv3.bias", "module.resBlock5.bn2.weight", "module.resBlock5.bn2.bias", "module.resBlock5.bn2.running_mean", "module.resBlock5.bn2.running_var", "module.resBlock5.conv4.weight", "module.resBlock5.conv4.bias", "module.resBlock5.bn3.weight", "module.resBlock5.bn3.bias", "module.resBlock5.bn3.running_mean", "module.resBlock5.bn3.running_var", "module.resBlock5.conv5.weight", "module.resBlock5.conv5.bias", "module.resBlock5.bn4.weight", "module.resBlock5.bn4.bias", "module.resBlock5.bn4.running_mean", "module.resBlock5.bn4.running_var", "module.upBlock1.conv1.weight", "module.upBlock1.conv1.bias", "module.upBlock1.bn1.weight", "module.upBlock1.bn1.bias", "module.upBlock1.bn1.running_mean", "module.upBlock1.bn1.running_var", "module.upBlock1.conv2.weight", "module.upBlock1.conv2.bias", "module.upBlock1.bn2.weight", "module.upBlock1.bn2.bias", "module.upBlock1.bn2.running_mean", "module.upBlock1.bn2.running_var", "module.upBlock1.conv3.weight", "module.upBlock1.conv3.bias", "module.upBlock1.bn3.weight", "module.upBlock1.bn3.bias", "module.upBlock1.bn3.running_mean", "module.upBlock1.bn3.running_var", "module.upBlock1.conv4.weight", "module.upBlock1.conv4.bias", "module.upBlock1.bn4.weight", "module.upBlock1.bn4.bias", "module.upBlock1.bn4.running_mean", "module.upBlock1.bn4.running_var", "module.upBlock2.conv1.weight", "module.upBlock2.conv1.bias", "module.upBlock2.bn1.weight", "module.upBlock2.bn1.bias", "module.upBlock2.bn1.running_mean", "module.upBlock2.bn1.running_var", "module.upBlock2.conv2.weight", "module.upBlock2.conv2.bias", "module.upBlock2.bn2.weight", "module.upBlock2.bn2.bias", "module.upBlock2.bn2.running_mean", "module.upBlock2.bn2.running_var", "module.upBlock2.conv3.weight", "module.upBlock2.conv3.bias", "module.upBlock2.bn3.weight", "module.upBlock2.bn3.bias", "module.upBlock2.bn3.running_mean", "module.upBlock2.bn3.running_var", "module.upBlock2.conv4.weight", "module.upBlock2.conv4.bias", "module.upBlock2.bn4.weight", "module.upBlock2.bn4.bias", "module.upBlock2.bn4.running_mean", "module.upBlock2.bn4.running_var", "module.upBlock3.conv1.weight", "module.upBlock3.conv1.bias", "module.upBlock3.bn1.weight", "module.upBlock3.bn1.bias", "module.upBlock3.bn1.running_mean", "module.upBlock3.bn1.running_var", "module.upBlock3.conv2.weight", "module.upBlock3.conv2.bias", "module.upBlock3.bn2.weight", "module.upBlock3.bn2.bias", "module.upBlock3.bn2.running_mean", "module.upBlock3.bn2.running_var", "module.upBlock3.conv3.weight", "module.upBlock3.conv3.bias", "module.upBlock3.bn3.weight", "module.upBlock3.bn3.bias", "module.upBlock3.bn3.running_mean", "module.upBlock3.bn3.running_var", "module.upBlock3.conv4.weight", "module.upBlock3.conv4.bias", "module.upBlock3.bn4.weight", "module.upBlock3.bn4.bias", "module.upBlock3.bn4.running_mean", "module.upBlock3.bn4.running_var", "module.upBlock4.conv1.weight", "module.upBlock4.conv1.bias", "module.upBlock4.bn1.weight", "module.upBlock4.bn1.bias", "module.upBlock4.bn1.running_mean", "module.upBlock4.bn1.running_var", "module.upBlock4.conv2.weight", "module.upBlock4.conv2.bias", "module.upBlock4.bn2.weight", "module.upBlock4.bn2.bias", "module.upBlock4.bn2.running_mean", "module.upBlock4.bn2.running_var", "module.upBlock4.conv3.weight", "module.upBlock4.conv3.bias", "module.upBlock4.bn3.weight", "module.upBlock4.bn3.bias", "module.upBlock4.bn3.running_mean", "module.upBlock4.bn3.running_var", "module.upBlock4.conv4.weight", "module.upBlock4.conv4.bias", "module.upBlock4.bn4.weight", "module.upBlock4.bn4.bias", "module.upBlock4.bn4.running_mean", "module.upBlock4.bn4.running_var", "module.logits.weight", "module.logits.bias". Unexpected key(s) in state_dict: "downCntx.conv1.weight", "downCntx.conv1.bias", "downCntx.conv2.weight", "downCntx.conv2.bias", "downCntx.bn1.weight", "downCntx.bn1.bias", "downCntx.bn1.running_mean", "downCntx.bn1.running_var", "downCntx.bn1.num_batches_tracked", "downCntx.conv3.weight", "downCntx.conv3.bias", "downCntx.bn2.weight", "downCntx.bn2.bias", "downCntx.bn2.running_mean", "downCntx.bn2.running_var", "downCntx.bn2.num_batches_tracked", "downCntx2.conv1.weight", "downCntx2.conv1.bias", "downCntx2.conv2.weight", "downCntx2.conv2.bias", "downCntx2.bn1.weight", "downCntx2.bn1.bias", "downCntx2.bn1.running_mean", "downCntx2.bn1.running_var", "downCntx2.bn1.num_batches_tracked", "downCntx2.conv3.weight", "downCntx2.conv3.bias", "downCntx2.bn2.weight", "downCntx2.bn2.bias", "downCntx2.bn2.running_mean", "downCntx2.bn2.running_var", "downCntx2.bn2.num_batches_tracked", "downCntx3.conv1.weight", "downCntx3.conv1.bias", "downCntx3.conv2.weight", "downCntx3.conv2.bias", "downCntx3.bn1.weight", "downCntx3.bn1.bias", "downCntx3.bn1.running_mean", "downCntx3.bn1.running_var", "downCntx3.bn1.num_batches_tracked", "downCntx3.conv3.weight", "downCntx3.conv3.bias", "downCntx3.bn2.weight", "downCntx3.bn2.bias", "downCntx3.bn2.running_mean", "downCntx3.bn2.running_var", "downCntx3.bn2.num_batches_tracked", "resBlock1.conv1.weight", "resBlock1.conv1.bias", "resBlock1.conv2.weight", "resBlock1.conv2.bias", "resBlock1.bn1.weight", "resBlock1.bn1.bias", "resBlock1.bn1.running_mean", "resBlock1.bn1.running_var", "resBlock1.bn1.num_batches_tracked", "resBlock1.conv3.weight", "resBlock1.conv3.bias", "resBlock1.bn2.weight", "resBlock1.bn2.bias", "resBlock1.bn2.running_mean", "resBlock1.bn2.running_var", "resBlock1.bn2.num_batches_tracked", "resBlock1.conv4.weight", "resBlock1.conv4.bias", "resBlock1.bn3.weight", "resBlock1.bn3.bias", "resBlock1.bn3.running_mean", "resBlock1.bn3.running_var", "resBlock1.bn3.num_batches_tracked", "resBlock1.conv5.weight", "resBlock1.conv5.bias", "resBlock1.bn4.weight", "resBlock1.bn4.bias", "resBlock1.bn4.running_mean", "resBlock1.bn4.running_var", "resBlock1.bn4.num_batches_tracked", "resBlock2.conv1.weight", "resBlock2.conv1.bias", "resBlock2.conv2.weight", "resBlock2.conv2.bias", "resBlock2.bn1.weight", "resBlock2.bn1.bias", "resBlock2.bn1.running_mean", "resBlock2.bn1.running_var", "resBlock2.bn1.num_batches_tracked", "resBlock2.conv3.weight", "resBlock2.conv3.bias", "resBlock2.bn2.weight", "resBlock2.bn2.bias", "resBlock2.bn2.running_mean", "resBlock2.bn2.running_var", "resBlock2.bn2.num_batches_tracked", "resBlock2.conv4.weight", "resBlock2.conv4.bias", "resBlock2.bn3.weight", "resBlock2.bn3.bias", "resBlock2.bn3.running_mean", "resBlock2.bn3.running_var", "resBlock2.bn3.num_batches_tracked", "resBlock2.conv5.weight", "resBlock2.conv5.bias", "resBlock2.bn4.weight", "resBlock2.bn4.bias", "resBlock2.bn4.running_mean", "resBlock2.bn4.running_var", "resBlock2.bn4.num_batches_tracked", "resBlock3.conv1.weight", "resBlock3.conv1.bias", "resBlock3.conv2.weight", "resBlock3.conv2.bias", "resBlock3.bn1.weight", "resBlock3.bn1.bias", "resBlock3.bn1.running_mean", "resBlock3.bn1.running_var", "resBlock3.bn1.num_batches_tracked", "resBlock3.conv3.weight", "resBlock3.conv3.bias", "resBlock3.bn2.weight", "resBlock3.bn2.bias", "resBlock3.bn2.running_mean", "resBlock3.bn2.running_var", "resBlock3.bn2.num_batches_tracked", "resBlock3.conv4.weight", "resBlock3.conv4.bias", "resBlock3.bn3.weight", "resBlock3.bn3.bias", "resBlock3.bn3.running_mean", "resBlock3.bn3.running_var", "resBlock3.bn3.num_batches_tracked", "resBlock3.conv5.weight", "resBlock3.conv5.bias", "resBlock3.bn4.weight", "resBlock3.bn4.bias", "resBlock3.bn4.running_mean", "resBlock3.bn4.running_var", "resBlock3.bn4.num_batches_tracked", "resBlock4.conv1.weight", "resBlock4.conv1.bias", "resBlock4.conv2.weight", "resBlock4.conv2.bias", "resBlock4.bn1.weight", "resBlock4.bn1.bias", "resBlock4.bn1.running_mean", "resBlock4.bn1.running_var", "resBlock4.bn1.num_batches_tracked", "resBlock4.conv3.weight", "resBlock4.conv3.bias", "resBlock4.bn2.weight", "resBlock4.bn2.bias", "resBlock4.bn2.running_mean", "resBlock4.bn2.running_var", "resBlock4.bn2.num_batches_tracked", "resBlock4.conv4.weight", "resBlock4.conv4.bias", "resBlock4.bn3.weight", "resBlock4.bn3.bias", "resBlock4.bn3.running_mean", "resBlock4.bn3.running_var", "resBlock4.bn3.num_batches_tracked", "resBlock4.conv5.weight", "resBlock4.conv5.bias", "resBlock4.bn4.weight", "resBlock4.bn4.bias", "resBlock4.bn4.running_mean", "resBlock4.bn4.running_var", "resBlock4.bn4.num_batches_tracked", "resBlock5.conv1.weight", "resBlock5.conv1.bias", "resBlock5.conv2.weight", "resBlock5.conv2.bias", "resBlock5.bn1.weight", "resBlock5.bn1.bias", "resBlock5.bn1.running_mean", "resBlock5.bn1.running_var", "resBlock5.bn1.num_batches_tracked", "resBlock5.conv3.weight", "resBlock5.conv3.bias", "resBlock5.bn2.weight", "resBlock5.bn2.bias", "resBlock5.bn2.running_mean", "resBlock5.bn2.running_var", "resBlock5.bn2.num_batches_tracked", "resBlock5.conv4.weight", "resBlock5.conv4.bias", "resBlock5.bn3.weight", "resBlock5.bn3.bias", "resBlock5.bn3.running_mean", "resBlock5.bn3.running_var", "resBlock5.bn3.num_batches_tracked", "resBlock5.conv5.weight", "resBlock5.conv5.bias", "resBlock5.bn4.weight", "resBlock5.bn4.bias", "resBlock5.bn4.running_mean", "resBlock5.bn4.running_var", "resBlock5.bn4.num_batches_tracked", "upBlock1.conv1.weight", "upBlock1.conv1.bias", "upBlock1.bn1.weight", "upBlock1.bn1.bias", "upBlock1.bn1.running_mean", "upBlock1.bn1.running_var", "upBlock1.bn1.num_batches_tracked", "upBlock1.conv2.weight", "upBlock1.conv2.bias", "upBlock1.bn2.weight", "upBlock1.bn2.bias", "upBlock1.bn2.running_mean", "upBlock1.bn2.running_var", "upBlock1.bn2.num_batches_tracked", "upBlock1.conv3.weight", "upBlock1.conv3.bias", "upBlock1.bn3.weight", "upBlock1.bn3.bias", "upBlock1.bn3.running_mean", "upBlock1.bn3.running_var", "upBlock1.bn3.num_batches_tracked", "upBlock1.conv4.weight", "upBlock1.conv4.bias", "upBlock1.bn4.weight", "upBlock1.bn4.bias", "upBlock1.bn4.running_mean", "upBlock1.bn4.running_var", "upBlock1.bn4.num_batches_tracked", "upBlock2.conv1.weight", "upBlock2.conv1.bias", "upBlock2.bn1.weight", "upBlock2.bn1.bias", "upBlock2.bn1.running_mean", "upBlock2.bn1.running_var", "upBlock2.bn1.num_batches_tracked", "upBlock2.conv2.weight", "upBlock2.conv2.bias", "upBlock2.bn2.weight", "upBlock2.bn2.bias", "upBlock2.bn2.running_mean", "upBlock2.bn2.running_var", "upBlock2.bn2.num_batches_tracked", "upBlock2.conv3.weight", "upBlock2.conv3.bias", "upBlock2.bn3.weight", "upBlock2.bn3.bias", "upBlock2.bn3.running_mean", "upBlock2.bn3.running_var", "upBlock2.bn3.num_batches_tracked", "upBlock2.conv4.weight", "upBlock2.conv4.bias", "upBlock2.bn4.weight", "upBlock2.bn4.bias", "upBlock2.bn4.running_mean", "upBlock2.bn4.running_var", "upBlock2.bn4.num_batches_tracked", "upBlock3.conv1.weight", "upBlock3.conv1.bias", "upBlock3.bn1.weight", "upBlock3.bn1.bias", "upBlock3.bn1.running_mean", "upBlock3.bn1.running_var", "upBlock3.bn1.num_batches_tracked", "upBlock3.conv2.weight", "upBlock3.conv2.bias", "upBlock3.bn2.weight", "upBlock3.bn2.bias", "upBlock3.bn2.running_mean", "upBlock3.bn2.running_var", "upBlock3.bn2.num_batches_tracked", "upBlock3.conv3.weight", "upBlock3.conv3.bias", "upBlock3.bn3.weight", "upBlock3.bn3.bias", "upBlock3.bn3.running_mean", "upBlock3.bn3.running_var", "upBlock3.bn3.num_batches_tracked", "upBlock3.conv4.weight", "upBlock3.conv4.bias", "upBlock3.bn4.weight", "upBlock3.bn4.bias", "upBlock3.bn4.running_mean", "upBlock3.bn4.running_var", "upBlock3.bn4.num_batches_tracked", "upBlock4.conv1.weight", "upBlock4.conv1.bias", "upBlock4.bn1.weight", "upBlock4.bn1.bias", "upBlock4.bn1.running_mean", "upBlock4.bn1.running_var", "upBlock4.bn1.num_batches_tracked", "upBlock4.conv2.weight", "upBlock4.conv2.bias", "upBlock4.bn2.weight", "upBlock4.bn2.bias", "upBlock4.bn2.running_mean", "upBlock4.bn2.running_var", "upBlock4.bn2.num_batches_tracked", "upBlock4.conv3.weight", "upBlock4.conv3.bias", "upBlock4.bn3.weight", "upBlock4.bn3.bias", "upBlock4.bn3.running_mean", "upBlock4.bn3.running_var", "upBlock4.bn3.num_batches_tracked", "upBlock4.conv4.weight", "upBlock4.conv4.bias", "upBlock4.bn4.weight", "upBlock4.bn4.bias", "upBlock4.bn4.running_mean", "upBlock4.bn4.running_var", "upBlock4.bn4.num_batches_tracked", "logits.weight", "logits.bias".

[MaxChanger]

Please provide more information such as how many cards are used for training.

This is a common general error, because adding module.* to the key prefix.

You could try like this: https://github.com/haomo-ai/MotionSeg3D/blob/dc4c95fcdba2f0819d2bbc4a419f231c55e9c6f3/modules/user.py#L63-L69

self.model = SalsaNextWithMotionAttention(self.parser.get_n_classes(), ARCH, num_batch=self.infer_batch_size)
self.model = nn.DataParallel(self.model)
checkpoint = "SalsaNextWithMotionAttention_refine_module_valid_best"
w_dict = torch.load(f"{self.modeldir}/{checkpoint}", map_location=lambda storage, loc: storage)
# self.model.load_state_dict(w_dict['main_state_dict'], strict=True)
self.model.load_state_dict({f"module.{k}":v for k,v in w_dict['main_state_dict'].items()}, strict=True)

[silence-tang]

Thanks a lot my friend. It seems like we just need to make a little change to the original code in user.py: line 64: self.model.load_state_dict(w_dict['state_dict'], strict=Ture) --> self.model.load_state_dict(w_dict['state_dict'], strict=False) line 70: self.model.load_state_dict(w_dict['state_dict'], strict=Ture) --> self.model.load_state_dict(w_dict['state_dict'], strict=False) I successfully solved this problem by making changes as shown above, and I hope it may help some other friends who also meet this problem.

[MaxChanger]

Hi, @silence-tang, I am very worried that if you use it like this, it will not load successfully. The logic of strict=False is to load values with equal keys. Although there is no error, it may not load any weight from the checkpoint. As shown above, I think use this command is better, because strict=False is used directly, other key mismatches cannot be avoided, It is very likely that it did not load any weights from the checkpoint.

# self.model.load_state_dict(w_dict['main_state_dict'], strict=True)
self.model.load_state_dict({f"module.{k}":v for k,v in w_dict['main_state_dict'].items()}, strict=True)

[silence-tang]

Thank you very much! @MaxChanger I just visualized the .label file generated by my method, but just as you said, it didn't load weight from the checkpoint and every point was labeled 9, which isn't what we expected. After applying the method you proposed, it finally worked in a proper way, thanks!