未验证 提交 73d78557 编写于 作者: F fengjiayi 提交者: GitHub

Fix a type error top_k_op (#5201)

* Fix Type error

* Fix error
上级 0f9858a1
...@@ -40,7 +40,7 @@ class TopkKernel : public framework::OpKernel<T> { ...@@ -40,7 +40,7 @@ class TopkKernel : public framework::OpKernel<T> {
const size_t k = static_cast<int>(ctx.Attr<int>("k")); const size_t k = static_cast<int>(ctx.Attr<int>("k"));
T* output_data = output->mutable_data<T>(ctx.GetPlace()); T* output_data = output->mutable_data<T>(ctx.GetPlace());
T* indices_data = indices->mutable_data<T>(ctx.GetPlace()); int64_t* indices_data = indices->mutable_data<int64_t>(ctx.GetPlace());
auto eg_input = EigenMatrix<T>::From(*input); auto eg_input = EigenMatrix<T>::From(*input);
...@@ -66,7 +66,7 @@ class TopkKernel : public framework::OpKernel<T> { ...@@ -66,7 +66,7 @@ class TopkKernel : public framework::OpKernel<T> {
}); });
for (size_t j = 0; j < k; j++) { for (size_t j = 0; j < k; j++) {
output_data[i * k + j] = vec[j].first; output_data[i * k + j] = vec[j].first;
indices_data[i * k + j] = vec[j].second; indices_data[i * k + j] = int64_t(vec[j].second);
} }
} }
} }
......
...@@ -9,7 +9,7 @@ class TestTopkOp(OpTest): ...@@ -9,7 +9,7 @@ class TestTopkOp(OpTest):
k = 1 k = 1
input = np.random.random((32, 84)).astype("float32") input = np.random.random((32, 84)).astype("float32")
output = np.ndarray((32, k)) output = np.ndarray((32, k))
indices = np.ndarray((32, k)) indices = np.ndarray((32, k)).astype("int64")
self.inputs = {'X': input} self.inputs = {'X': input}
self.attrs = {'k': k} self.attrs = {'k': k}
...@@ -32,7 +32,7 @@ class TestTopkOp3d(OpTest): ...@@ -32,7 +32,7 @@ class TestTopkOp3d(OpTest):
input = np.random.random((32, 2, 84)).astype("float32") input = np.random.random((32, 2, 84)).astype("float32")
input_flat_2d = input.reshape(64, 84) input_flat_2d = input.reshape(64, 84)
output = np.ndarray((64, k)) output = np.ndarray((64, k))
indices = np.ndarray((64, k)).astype("int") indices = np.ndarray((64, k)).astype("int64")
# FIXME: should use 'X': input for a 3d input # FIXME: should use 'X': input for a 3d input
self.inputs = {'X': input_flat_2d} self.inputs = {'X': input_flat_2d}
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册