Image Tensors Return As Zero When num_workers > 0

I am facing an issue with multiprocessing. I am trying to load my .pt data as dataloaders. Everything works fine when I set the num_workers = 0. But when I set it to a value greater than 0, the tensor values become zero.

Below is my code to load the data. The shape of the tensors are all ok. Just the range of the images/tensors are [0, 0]. Would you please suggest what's wrong here? Thank you in advance.

class SRDataset(Dataset):
    def __init__(self, hr_tensors, lr_tensors, transform=None):
        self.hr_tensors = hr_tensors
        self.lr_tensors = lr_tensors
        self.transform = transform

        assert len(self.hr_tensors) == len(self.lr_tensors), \
            "Number of HR and LR images must be equal"

    def __len__(self):
        return len(self.hr_tensors)

    def __getitem__(self, idx):
        hr_img = self.hr_tensors[idx]
        lr_img = self.lr_tensors[idx]

        if self.transform:
            hr_img = self.transform(hr_img)
            lr_img = self.transform(lr_img)

        return {'lr': lr_img, 'hr': hr_img}

def create_sr_datasets_and_loaders(hr_tensors, lr_tensors, batch_size=32, transform=None):

    full_dataset = SRDataset(hr_tensors, lr_tensors, transform)

    total_size = len(full_dataset)
    train_size = int(0.8 * total_size)
    val_size = int(0.1 * total_size)
    test_size = total_size - train_size - val_size

    train_dataset, val_dataset, test_dataset = random_split(
        full_dataset,
        [train_size, val_size, test_size],
        generator=torch.Generator().manual_seed(seed)
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=1,
        drop_last=True,
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=1,
        drop_last=True,
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=1,
        drop_last=True,
    )

    return {
        'datasets': {
            'train': train_dataset,
            'val': val_dataset,
            'test': test_dataset
        },
        'loaders': {
            'train': train_loader,
            'val': val_loader,
            'test': test_loader
        }
    }

data = create_sr_datasets_and_loaders(
    hr_tensors,
    lr_tensors,
    batch_size=64,
    transform=None
)

train_loader = data['loaders']['train']
val_loader = data['loaders']['val']
test_loader = data['loaders']['test']

print(f"Total images: {len(hr_tensors)}")
print(f"Training set: {len(data['datasets']['train'])} images")
print(f"Validation set: {len(data['datasets']['val'])} images")
print(f"Test set: {len(data['datasets']['test'])} images")

train_batch = next(iter(train_loader))
print(f"\nTrain batch LR shape: {train_batch['lr'].shape}")
print(f"Train batch HR shape: {train_batch['hr'].shape}")

lr_sample = train_batch['lr'][0]

hr_sample = train_batch['hr'][0]

plt.figure(figsize=(12, 6))

plt.subplot(1, 2, 1)

plt.imshow(lr_sample.squeeze().cpu().numpy(), cmap='gray')

plt.title(f"Low Resolution Image (Shape: {lr_sample.shape})")

plt.axis('off')

plt.subplot(1, 2, 2)

plt.imshow(hr_sample.squeeze().cpu().numpy(), cmap='gray')

plt.title(f"High Resolution Image (Shape: {hr_sample.shape})")

plt.axis('off')

plt.tight_layout()

plt.show()

print(f"LR Image - Min: {lr_sample.min().item():.4f}, Max: {lr_sample.max().item():.4f}")

print(f"HR Image - Min: {hr_sample.min().item():.4f}, Max: {hr_sample.max().item():.4f}")

Initially, it was having some problems loading the data when I set the num_workers > 0. Then LLMs suggested me to add the following lines of code and then it ran successfully. But still the images are all black.

import dill

import pickle

import torch.multiprocessing as mp

import multiprocessing.reduction as reduction

reduction.ForkingPickler.dumps = dill.dumps

mp.set_start_method("spawn", force=True)

torch.multiprocessing.set_sharing_strategy("file_system")

pickle.HIGHEST_PROTOCOL = 5

For your information, I am using Ubuntu 22.04.5 LTS WSL on Windows 11 and I am testing the code on a jupyter notebook, not a python file.