Remove override now that autograph has been fixed.
@@ -56,11 +56,11 @@ class LinearRegression(model.Model):
def local_variables(self):
return [self._num_examples, self._num_batches, self._loss_sum]
- @tf.contrib.eager.defun(autograph=False)
+ @tf.contrib.eager.defun
def _predict(self, x):
return tf.matmul(x, self._a) + self._b + self._c
- @tf.contrib.eager.defun(autograph=False)
+ @tf.contrib.eager.defun
def forward_pass(self, batch, training=True):
del training # Unused
@@ -76,7 +76,7 @@ class LinearRegression(model.Model):
average_loss = total_loss / tf.cast(num_examples, tf.float32)
return model.BatchOutput(loss=average_loss, predictions=predictions)
- @tf.contrib.eager.defun(autograph=False)
+ @tf.contrib.eager.defun
def aggregated_outputs(self):
return collections.OrderedDict(
[('num_examples', self._num_examples), ('num_batches',
@@ -92,7 +92,7 @@ class LinearRegression(model.Model):
class TrainableLinearRegression(LinearRegression, model.TrainableModel):
"""A LinearRegression with trivial SGD training."""
- @tf.contrib.eager.defun(autograph=False)
+ @tf.contrib.eager.defun
def train_on_batch(self, batch):
# Most users won't implement this, and let us provide the optimizer.
fp = self.forward_pass(batch)
