Update main.py

main.py

@@ -76,7 +76,7 @@ def run_main(
     input_folder = os.path.join(base_input_folder, name)
     mask_list = []
     for mask_np in mask_np_list:
-        mask = torch.from_numpy(mask_np)
+        mask = torch.from_numpy(mask_np.astype(np.uint8))
         mask_list.append(mask)
     
     #mask_list, mask_label_list = load_mask(input_folder)