import tensorflow as tf

                            aux_logits = self._add_aux_head(*layers[-1], K, is_train)
                        aux_logits_list.append(aux_logits)

            # Global average pooling
            (X, w, h, ch) = layers[-1]
            X = self._add_global_avg_pool(X, w, h, ch)

            # Add dropout if training
            if is_train:
                X = tf.nn.dropout(X,  dropout_keep_prob)

            # Compute logits from X
            with tf.variable_scope('fully_connected'):
                logits = self._add_fully_connected(X, (ch,), K)

        return (logits, aux_logits_list)

    def _optimize(self, loss, step, **knobs):
        opt_momentum = knobs['opt_momentum'] # Momentum optimizer momentum
        grad_clip_norm = knobs['grad_clip_norm'] # L2 norm to clip gradients by

        # Compute learning rate, gradients
        tf_trainable_vars = tf.trainable_variables()
        lr = self._get_learning_rate(step, **knobs)
