Dataloader with sub-folder

I have a dataset that the training data are split into 4 folders with the same size of the data in each (3000 data). I wanted to create a data loader that loads only 3000 data but instead of sampling from one folder, it randomly samples from one of the sub-folders. So, I have a Dataset object that in its getitem function, I have:

 def __getitem__(self, idx):
        iq = int(torch.randint(0, len(self.img_address), (1,)))
        img_address_ = self.img_address[iq]
        ref_address_ = self.ref_address[iq]

        img_address = img_address_[idx]
        ref_address = ref_address_[idx]
        img = self.loader(img_address)
        ref = self.loader(ref_address)
       return img, ref
def __len__(self):
        return 3000

However, when I sample from the data loader, I get 4 samples (number of sub-folders) instead of batch size! I tested the dataset.getitem(N), and it works fine. The problem starts when I wrap the dataset in a Dataloader! I changed the number of workers, batch size, shuffle, and … but neither of them worked!

Can you please help me to find where I am doing wrong?

Your code works for me using random data:

class MyDataset(torch.utils.data.Dataset):
    def __init__(self):
        self.datasets = [torch.arange(i+i*100, i+i*100+100) for i in range(4)]

    def __getitem__(self, idx):
        iq = int(torch.randint(0, len(self.datasets), (1,)))
        dataset = self.datasets[iq]

        data = dataset[idx]
        return data
   
    def __len__(self):
        return len(self.datasets[0])

dataset = MyDataset()
print(dataset.datasets)

loader = torch.utils.data.DataLoader(dataset, batch_size=10)
for data in loader:
    print(data)

As you can see the batch_size is used and the samples are drawn randomly from the internal datasets.

Thank you for your help. I am having this problem when I set pin_memory=True on my data loader! I am not sure why this is happening!

Are you seeing the same issue using my code snippet with pin_memory=True?

No, I actually found a bug in my code that caused the problem! Sorry for bothering you and do appreciate your prompt answers!