The number of labels during training and testing are different

[HripsimeS]

Hello. I would like to mention one important mark concerning the visualization of labels. I have totally 5 labels (including a background) and after training process I can see all 4 labels in true_masks and predicted_masks. I used following Compose during training:
post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=5)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=5)])

But during evaluation of model on original images/masks visualization shows only 1 label for both true_masks and predicted_masks.
I assume the issue comes from post_transform, although I mentioned to_onehot=5

post_transforms = Compose([
EnsureTyped(keys="pred"),
Invertd(keys="pred", transform=test_transforms, orig_keys="image", meta_keys="pred_meta_dict",
orig_meta_keys="image_meta_dict", meta_key_postfix="meta_dict", nearest_interp=False,to_tensor=True,),
AsDiscreted(keys="pred", argmax=True, to_onehot=5),
AsDiscreted(keys="mask", to_onehot=5),
])

Do you know by chance what could cause these changes with number of labels. Thanks in advance for your help!

[HripsimeS]

This ticket has been already answered on discussions for tutorials channel, therefore can be closed. Thanks :)
Project-MONAI/tutorials#609