I created a dataset that loads a single data sample at a time on demand (1 sample consists of multiple images), and I have a data loader with a small batch size. When I try to show just the first few batches of my dataset, the loader keeps trying to iterate through my entire dataset instead of pulling out just the few data samples:
class FaceDataset(Dataset):
    def __init__(self):
        df = pd.read_csv("data/positions.csv")
        df["filename"] = df["id"].astype("str") + ".jpg"

        self.filenames = df["filename"].tolist()
        self.targets = torch.FloatTensor(list(zip(df["x"], df["y"])))
        self.head_angle = torch.FloatTensor(df['head_angle'].tolist())

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

    def __getitem__(self, idx):
        sample = {
            "targets": self.targets[idx],
            "head_angle": self.head_angle[idx]
        }

        for img_type in ["face", "face_aligned", "l_eye", "r_eye", "head_pos"]:
            img = Image.open("data/{}/{}".format(img_type, self.filenames[idx]))
            img = torch.from_numpy(np.array(img))
            # img /= 255
            sample[img_type] = img

        return sample

ds = FaceDataset()
data = DataLoader(ds, batch_size=2, shuffle=True, num_workers=2)

for i_batch, sample_batched in enumerate(data):
    print(i_batch, sample_batched)

    if i_batch == 1:
        break

The length of my dataset is about 15k, and the loader seems to be trying to load everything instead of just a single batch of 2. Jupyter seems to just freeze before anything actually gets printed. Am I creating my dataset/loader incorrectly?
When I try to iterate over the dataset, I seem to get the correct data back for 1 training sample. The problem only seems to come in when iterating over the data loader object:
for batch in ds:
    print(batch)
    break

EDIT: I think Ive narrowed this down to the num_workers param. When I set it to 0 then the whole thing works as expected, but when I set it to > 0 then nothing gets printed?