We will use pytorch Dataset Class to create our dataset

dataset_creation.py

from torch.utils.data.dataset import Dataset
from torchvision.transforms import Compose, ToTensor
class MaskDataset(Dataset):
        """ Masked faces dataset
        0 = 'no mask'
        1 = 'mask'
        """
        def __init__(self, train_data):
            self.train_data = train_data
            self.transformations = Compose([
                ToTensor(), # [0, 1]
            ])
        
        def __getitem__(self, key):
            if isinstance(key, slice):
                raise NotImplementedError('slicing is not supported')                    
            return [
                self.transformations(self.train_data[key][0]),
                torch.tensor(self.train_data[key][1]) # pylint: disable=not-callable
            ]
        
        def __len__(self):
            return len(self.train_data)