The execution of SimpleITK's STAPLE algorithm gets stuck in certain cases

Question

I am running the STAPLE algorithm to ensemble 9 segmentations from several models (each segmentation has 3 labels (1,2,3) and the background (0)), but it is getting stuck in some cases (specifically 2 cases). I wonder if someone is having the same issue and if it is related to the fact that not all segmentations have all labels, e.g., some segmentations only have label 3.

The code is as follows:

def create_STAPLE_segmentation(segmentations_arg, save_path, file_name):
    seg_mask_stack = []
    for seg1 in segmentations_arg:
        # Creating masks to fill the holes created by having multilabel
        seg1_mask = np.where(seg1 >= 1, 1, 0).astype(np.int16)
        seg1_mask_sitk = sitk.GetImageFromArray(seg1_mask) 
        seg_mask_stack.append(seg1_mask_sitk)

    # Run STAPLE algorithm
    STAPLE_seg_mask_sitk = sitk.STAPLE(seg_mask_stack, 1)
    # Mask created to fill the holes!
    STAPLE_seg_mask = sitk.GetArrayFromImage(STAPLE_seg_mask_sitk)

    # Doing STAPLE for each label
    # List of foreground values for each segmentation (assuming labels are 1, 2, and 3)
    foreground_values = [1, 2, 3]
    new_seg_list = []
    for foreground_value in foreground_values:
        seg_stack = []
        for seg1 in segmentations_arg:
            new_seg1 = np.where(seg1 == foreground_value, 1, 0).astype(np.int16)# STAPLE requires we cast into int16 arrays
            seg1_sitk = sitk.GetImageFromArray(new_seg1) 
            seg_stack.append(seg1_sitk)
        # Run STAPLE algorithm
        STAPLE_seg_sitk = sitk.STAPLE(seg_stack, 1)
        # Convert back to numpy array
        STAPLE_seg = sitk.GetArrayFromImage(STAPLE_seg_sitk)
        new_seg_list.append(STAPLE_seg)


    STAPLE_seg_all = np.zeros_like(STAPLE_seg_mask)
    STAPLE_seg_all[STAPLE_seg_mask>0.5] = 1
    for index, sub_lab in enumerate(new_seg_list):
        STAPLE_seg_all[sub_lab>0.5] = index+1 

    # Save the NIfTI image to a .nii.gz file
    output_file = join(save_path, file_name.split('/')[-1])
    nifti_for_metadata = nib.load(file_name)
    nifti_img = nib.Nifti1Image(STAPLE_seg_all, affine=nifti_for_metadata.affine, header=nifti_for_metadata.header)
    nib.save(nifti_img, output_file)
    #return STAPLE_seg_all

Thank you in advance!

I was expecting to have the segmentation ensembled by the STAPLE algorithm, as I have for the other cases.