The problem is solved according to @OldBoy input to move the increment.

full code

mask_i = 0
a = 0

f = open("masking_log.txt", "w")

for im in large_image_stack_512:
    image_success_flag = 0
    mask_i= 0
    
    while(image_success_flag < 1):
        jpeg_im = Image.open(os.path.join(ROOT_DIR,im))
        print(os.path.join(ROOT_DIR,im))
        # Semantic segmentation
        segmentation = semantic_segmentation_nvidia(jpeg_im)
        print("the length of current segmentation labels are: ", len(segmentation))
        while(mask_i < len(segmentation)):
           image_mask = segmentation[mask_i]["label"]
           print(image_mask)
           if(image_mask == "water"):
               print("correct mask")
               water_mask = segmentation[mask_i]["mask"]
               imar = np.asarray(water_mask)
               plt.imsave('D:/semester_12/Data_Irredeano/'+'img_'+str(a)+'.jpg', imar, cmap="gray") 
               print("here")
               f.write("image " + str(a) + "\nsuccess-label at " + str(mask_i) + "\nwith dir: " + str(im)  + "\n with mask labeled as: " + str(water_mask_label) + '\n\n')
               print("mask-succesfully saved")
               mask_i = 0 
               break
           elif(image_mask != "water"):
               mask_i+=1
               print("number of mask: ", mask_i)
               if(mask_i == len(segmentation)):
                    print("this image has no correct mask, check later")
                    f.write("image " + str(a) + "\n unsuccess-labelling (has no 'water' label) final \n mask_i value: " + str(mask_i) + "\nwith dir: " + str(im)  + "\n check later " + '\n\n')
               image_success_flag=+1  
        a+=1


f.close()

Basically instead of selecting mask by checking the segmentation[mask_i]["label] I check if the 'cursor' or mask_i if it's smaller than the length of the List (len(segmentation)). The += also contribute to the problem since it's adding first then change the number, as implied here and here, because of that my 'cursor variable' can move beyond the array size before checking the segmentation[mask_i]["label]. But I don't think I have other choice other than that to increment since =+ is just redefining itself.

Other than that I also add another while condition to make sure the code runs while the mask_i is below the size of the List, so the program become "if it's still below the list size, check whether mask is "water" or not "water".

Although the program is finished and can mask most of the images, I still have to log several different images since not all of them had "water" as the label, but something close like "Sea" that we humans can intuitively say they are basically the same but computer that can only compare strings not.

Thank you again for everyone who's willing to help, I'm accepting if there's better way to do it