monai data loader stops after few iteration, NVIDIA RTX 4090 GPU

[YazdanSalimi]

I define data_loader_train by this transform.
transform = monai.transforms.Compose(
[
monai.transforms.LoadImaged(keys=["input_image", "output_image"],
ensure_channel_first = True, image_only = False),
monai.transforms.Spacingd(keys=["input_image", "output_image"], pixdim = pixdim,
mode = [input_interpolation, output_interpolation]),
monai.transforms.EnsureTyped(keys=["input_image", "output_image"]),
monai.transforms.Orientationd(keys=["input_image", "output_image"], axcodes=orientation),
monai.transforms.CropForegroundd(keys=["input_image", "output_image"], source_key="input_image", margin = ImageCropForeGroundMargin),

                        monai.transforms.ScaleIntensityRanged(keys = "output_image", 
                                                              a_min = output_lower_intensity, a_max = output_upper_intensity,
                                                              clip=clip, b_min = output_b_min, b_max = output_b_max) if not scale_image_to_self_max else monai.transforms.ScaleIntensityd(keys = ["output_image"],
                                                                                                    minv = output_lower_intensity, maxv=output_upper_intensity),
                        monai.transforms.ScaleIntensityRanged(keys = "input_image", a_min = input_lower_intensity,
                                                              b_min = input_b_min, b_max = input_b_max),
                        monai.transforms.RandCropByPosNegLabeld(keys=["input_image", "output_image"],
                                                                label_key = "output_image", 
                                                                spatial_size = patch_size,
                                                                allow_smaller = True, 
                                                                image_key= "input_image", num_samples=num_samples,
                                                                pos = positive, neg = negative),
                        monai.transforms.ResizeWithPadOrCropd(keys=["input_image", "output_image"],
                                                              spatial_size= patch_size),
                        monai.transforms.RandFlipd(keys=["input_image", "output_image"]), # add augmentations after all processes
                        
                    ]

then try to load the batch to GPU in a loop for testing the dataloader. as below:

for repearted_evaluation in range(data_validity_epochs):
for batch_data in tqdm(data_loader_train,):
input_image, output_image = (torch.squeeze(batch_data["input_image"], dim = -1).to(device),
torch.squeeze(batch_data["output_image"], dim = -1).to(device))

after few epochs, simetimes 100, sometimes five epoochs this loop stops without any error.
I have tried it ttwith cuda 12.1 and 11.8 and diffrerenet number of workers from 0 to 10. I am using windows, CPU corei9 13900KF and GPU RTX 4090. erro is still there.

the exact same code and data work fine on another device with RTX 3080 and CPU corei712700KF.
Could you please help me to solve it?
Thanks

[YazdanSalimi]

As an update if I change "torch.squeeze(batch_data["output_image"], dim = -1).to(device))" to batch_data["output_image"] the error does not happen. it is about torch.squeeze?

[KumoLiu]

Hi @YazdanSalimi, thanks for your interest here.
It's not easy to identify what the error is, maybe you can try running with num_workers = 0 to get a proper debug trace.

[YazdanSalimi]

Thank you for your response,
Unfortunately it happens even with num_workers = 0.