import tensorflow as tf

        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)
        grads = tf.gradients(loss, tf_trainable_vars)
        self._mark_for_monitoring('lr', lr)

        # Clip gradients
        if grad_clip_norm > 0:
            grads = [tf.clip_by_norm(x, grad_clip_norm) for x in grads]

        # Init optimizer
        opt = tf.train.MomentumOptimizer(lr, opt_momentum, use_locking=True, use_nesterov=True)
        train_op = opt.apply_gradients(zip(grads, tf_trainable_vars), global_step=step)

        return train_op

    def _preprocess(self, images, classes, is_train=False, **knobs):
        batch_size = knobs['batch_size']
        cutout_size = knobs['cutout_size']
