def _load(self, dataset='train'):
        """Load the data in memory.
        Args:
            dataset: string in ['train', 'test']
        """
        data, labels = None, None
        if dataset is 'train':
            files = [os.path.join(self.cifar10_dir, 'data_batch_%d' % i) for i in range(1, 6)]
        else:
            files = [os.path.join(self.cifar10_dir, 'test_batch')]

        for file in files:
            if not os.path.exists(file):
                raise FileNotFoundError('Failed to find file: ' + file)

        # Load the data from the batch files
        for file in files:
            with open(file, 'rb') as f:
                cifar10 = pickle.load(f, encoding='latin1')

            if labels is None:
                labels = np.array(cifar10['labels'])
            else:
                labels = np.concatenate((labels, cifar10['labels']), axis=0)

            if data is None:
                data = cifar10['data']
            else:
                data = np.concatenate((data, cifar10['data']), axis=0)

        # Adapt the format of the data to our convnet
        data = np.array(data, dtype=float) / 255.0
        data = data.reshape([-1, self.num_channels, self.img_size, self.img_size])
        data = data.transpose([0, 2, 3, 1])

        # One-hot encode labels (see https://stackoverflow.com/a/42874726)
        labels = np.eye(self.num_classes)[np.array(labels).reshape(-1)]

        if dataset is 'train':
            self._train_data, self._train_labels = data, labels
        else:
            self._test_data, self._test_labels = data, labels