from tensorflow.contrib.eager.python.examples.linear_regression import linear_regression

        batch_size=batch_size,
        num_batches=num_batches)
    burn_in_dataset = dataset.take(10)

    model = linear_regression.LinearModel()

    with tf.device(device()):
      optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.1)
