Merge pull request #13101 from zenghsh3/develop

Fix bug of sampling_id op

Author: seiriosPlus

paddle/fluid/operators/sampling_id_op.h

@@ -53,7 +53,7 @@ class SamplingIdKernel : public framework::OpKernel<T> {
         static_cast<T>(context.Attr<float>("min")),
         static_cast<T>(context.Attr<float>("max")));
 
-    std::vector<T> ids(batch_size);
+    std::vector<int64_t> ids(batch_size);
     for (int i = 0; i < batch_size; ++i) {
       T r = dist(engine);
       int idx = width - 1;
@@ -63,7 +63,7 @@ class SamplingIdKernel : public framework::OpKernel<T> {
           break;
         }
       }
-      ids[i] = ins_vector[idx];
+      ids[i] = int64_t(idx);
     }
 
     std::vector<int64_t> out_dim;

python/paddle/fluid/tests/unittests/test_sampling_id_op.py

@@ -25,9 +25,9 @@ def setUp(self):
         self.op_type = "sampling_id"
         self.use_mkldnn = False
         self.init_kernel_type()
-        self.X = np.random.random((8, 4)).astype('float32')
+        self.X = np.random.random((100, 10)).astype('float32')
         self.inputs = {"X": self.X}
-        self.Y = np.random.random(8).astype('float32')
+        self.Y = np.random.random(100).astype('int64')
         self.outputs = {'Out': self.Y}
         self.attrs = {'max': 1.0, 'min': 0.0, 'seed': 1}
 
@@ -36,6 +36,16 @@ def test_check_output(self):
         y1 = self.out
         self.check_output_customized(self.verify_output)
         y2 = self.out
+
+        # check dtype
+        assert y1.dtype == np.int64
+        assert y2.dtype == np.int64
+
+        # check output is index ids of inputs
+        inputs_ids = np.arange(self.X.shape[1])
+        assert np.isin(y1, inputs_ids).all()
+        assert np.isin(y2, inputs_ids).all()
+
         self.assertTrue(np.array_equal(y1, y2))
         self.assertEqual(len(y1), len(self.Y))