aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/python/kernel_tests/in_topk_op_test.py
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/python/kernel_tests/in_topk_op_test.py')
-rw-r--r--tensorflow/python/kernel_tests/in_topk_op_test.py5
1 files changed, 1 insertions, 4 deletions
diff --git a/tensorflow/python/kernel_tests/in_topk_op_test.py b/tensorflow/python/kernel_tests/in_topk_op_test.py
index 37e9a8e3d1..fafeea8ec0 100644
--- a/tensorflow/python/kernel_tests/in_topk_op_test.py
+++ b/tensorflow/python/kernel_tests/in_topk_op_test.py
@@ -22,7 +22,6 @@ import numpy as np
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors_impl
-from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.platform import test
@@ -77,9 +76,7 @@ class InTopKTest(test.TestCase):
k = constant_op.constant(3)
np_ans = np.array([False, True])
with self.test_session():
- # TODO (yongtang): The test will be switch to nn_ops.in_top
- # once nn_ops.in_top points to _in_top_kv2 later
- precision = gen_nn_ops._in_top_kv2(predictions, target, k)
+ precision = nn_ops.in_top_k(predictions, target, k)
out = precision.eval()
self.assertAllClose(np_ans, out)
self.assertShapeEqual(np_ans, precision)