I have the DataLoader class for 1 image and labels:
class CoronalDataset(Dataset):
    def __init__(self, pet_type='interim', train=True, transform=None):
        train_valid = 'train' if train else 'valid'
        # Create the Tensor from the numpy array 
        self.x = torch.from_numpy(np.load(f'{DATASET_DIR}/{pet_type}/rgb/x_{train_valid}_coronal.npy')).permute(0, 3, 1, 2).type(torch.FloatTensor)
        print("Shape of X: ", self.x.shape)
        print(self.x.dtype)
        
        self.y = torch.from_numpy(np.load(f'{DATASET_DIR}/{pet_type}/rgb/y_{train_valid}_123_45.npy'))
        print("Shape of Y:", self.y.unique(return_counts=True))
        labellist = self.y.tolist()
        result = sorted([(x, labellist.count(x)) for x in [[1.0, 0.0],[0.0, 1.0]]], key=lambda y: y[1])
        for elem in result:
            print('{} {}'.format(elem[0], elem[1]))

        print(self.y.dtype)
        
        self.len = len(self.x)
        self.transform = transform
        
    def __getitem__(self, index):
        sample = self.x[index], self.y[index]

        if self.transform:
            sample = self.transform(sample1)

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