loading an image using the image path
data_path = 'data/train/'
train_dataset = torchvision.datasets.ImageFolder(
      root=data_path,
      transform=torchvision.transforms.ToTensor()
    )

get the images into train_loader using DataLoader, you can get single image by changing the batch_size to 1, we use shuffle True because the model we train should be robust to avoid overfitting
train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=64,
        shuffle=True
    )