import tensorflow as tf


        return (images_batch, classes_batch, dataset_init_op)

    def _get_drop_path_keep_prob(self, layers_ratio, step, is_train=False, **knobs):
        batch_size = knobs['batch_size']
        drop_path_keep_prob = knobs['drop_path_keep_prob'] # Base keep prob for drop path
        drop_path_decay_epochs = knobs['drop_path_decay_epochs']
        N = self._train_params['N']

        # Only drop path during training
        keep_prob = tf.constant(1, dtype=tf.float32)
        if is_train:
            # Decrease keep prob deeper into network
            keep_prob = 1 - layers_ratio * (1 - drop_path_keep_prob)

            # Decrease keep prob with increasing steps
            steps_per_epoch = math.ceil(N / batch_size)
            steps_ratio = tf.minimum(((step + 1) / steps_per_epoch) / drop_path_decay_epochs, 1)
            keep_prob = 1 - steps_ratio * (1 - keep_prob)
            keep_prob = tf.cast(keep_prob, tf.float32)
