未验证 提交 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> {
const size_t k = static_cast<int>(ctx.Attr<int>("k"));
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);
......@@ -66,7 +66,7 @@ class TopkKernel : public framework::OpKernel<T> {
});
for (size_t j = 0; j < k; j++) {
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):
k = 1
input = np.random.random((32, 84)).astype("float32")
output = np.ndarray((32, k))
indices = np.ndarray((32, k))
indices = np.ndarray((32, k)).astype("int64")
self.inputs = {'X': input}
self.attrs = {'k': k}
......@@ -32,7 +32,7 @@ class TestTopkOp3d(OpTest):
input = np.random.random((32, 2, 84)).astype("float32")
input_flat_2d = input.reshape(64, 84)
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
self.inputs = {'X': input_flat_2d}
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册