import tensorflow as tf



class StudentTask(base_model.BaseTask):

  @base_layer.initializer
  def __init__(self, params):
    super(StudentTask, self).__init__(params)
    p = self.params
    with tf.variable_scope(p.name):
      self.CreateVariable('x',
                          py_utils.WeightParams(
                              shape=[], init=py_utils.WeightInit.Uniform()))

  def ComputePredictions(self, theta, input_batch):
    return theta.x

