Update tpu_estimator.py

Fixing zihangdai#43. This does not affect the current repo because it is not used.
gtkafka · Jun 24, 2019 · b7a748a · b7a748a
1 parent 331dd00
commit b7a748a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tpu_estimator.py b/tpu_estimator.py
@@ -2899,7 +2899,7 @@ def multi_tpu_eval_steps_on_single_shard():
       outputs_from_all_shards=False,
       device_assignment=ctx.device_assignment)
 
-  loss = rte[0]
+  loss = ret[0]
   scaffold = _get_scaffold(captured_scaffold_fn)
   return compile_op, loss, host_calls, scaffold, captured_eval_hooks.get()