import tensorflow as tf

        def _preprocess_train(image, clazz):
            # Do random crop + horizontal flip for each train image
            image = tf.pad(image, [[4, 4], [4, 4], [0, 0]])
            image = tf.image.random_crop(image, (w, h, in_ch))
            image = tf.image.random_flip_left_right(image)

            if cutout_size > 0:
                image = self._do_cutout(image, w, h, cutout_size)

            return (image, clazz)

        (images, classes) = _prepare(images, classes)
        dataset = tf.data.Dataset.from_tensor_slices((images, classes)).repeat()
        if is_train:
            dataset = dataset.apply(tf.data.experimental.map_and_batch(map_func=_preprocess_train, batch_size=batch_size))
        else:
            dataset = dataset.batch(batch_size)
        dataset_itr = dataset.make_initializable_iterator()
        (images_batch, classes_batch) = dataset_itr.get_next()
        dataset_init_op = dataset_itr.initializer

        return (images_batch, classes_batch, dataset_init_op)

    def _get_drop_path_keep_prob(self, layers_ratio, step, is_train=False, **knobs):
