If anyone is looking for a different method, torchvision has a utility to accomplish this task cleanly.
https://pytorch.org/docs/stable/torchvision/datasets.html#torchvision.datasets.ImageFolder 199
There’s a small caveat with this though, images in a particular batch are not guaranteed to come from different classes.
It assumes the images are arranged in the following manner:
root_dir/
    class_1/
        001.png
        002.png
        ...
    class_2/
        001.png
        002.png
        ...
    ...

from torchvision import datasets, transforms
from torch.utils import data

dataset = datasets.ImageFolder(root = root_dir, 
                transform = transforms.ToTensor())

loader = data.DataLoader(dataset, batch_size = 8, shuffle = True)