Fix a type error top_k_op (#5201)

* Fix Type error * Fix error
PaddlePaddle · Oct 30, 2017 · 73d7855 · 73d7855
1 parent 0f9858a
commit 73d7855
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/paddle/operators/top_k_op.h b/paddle/operators/top_k_op.h
@@ -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);
       }
     }
   }

diff --git a/python/paddle/v2/framework/tests/test_top_k_op.py b/python/paddle/v2/framework/tests/test_top_k_op.py
@@ -9,7 +9,7 @@ def setUp(self):
         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 @@ def setUp(self):
         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}