From eda5d0f41f125020c14e9f77f8a47d68e8f33fb1 Mon Sep 17 00:00:00 2001
From: ooooo <3164076421@qq.com>
Date: Thu, 1 Aug 2024 21:54:29 +0800
Subject: [PATCH] update place

---
 paddle/phi/infermeta/binary.cc            |  3 +-
 test/legacy_test/test_shuffle_batch_op.py | 46 +++++++++++++----------
 2 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
index 7b181159a65ffa..6877f707e80a91 100644
--- a/paddle/phi/infermeta/binary.cc
+++ b/paddle/phi/infermeta/binary.cc
@@ -3760,9 +3760,8 @@ void ShuffleBatchInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
   seed_out->share_dims(seed);
   seed_out->share_lod(seed);
-  seed_out->set_dtype(x.dtype());
+  seed_out->set_dtype(seed.dtype());
   shuffle_idx->set_dims(phi::make_ddim({-1}));
-  shuffle_idx->set_dtype(x.dtype());
 }
 
 void SequenceMaskInferMeta(const MetaTensor& x,
diff --git a/test/legacy_test/test_shuffle_batch_op.py b/test/legacy_test/test_shuffle_batch_op.py
index 8015580e77160b..bf508065d666dc 100644
--- a/test/legacy_test/test_shuffle_batch_op.py
+++ b/test/legacy_test/test_shuffle_batch_op.py
@@ -89,13 +89,16 @@ def get_shape(self):
 
 class TestShuffleBatchAPI(unittest.TestCase):
     def setUp(self):
+        self.places = [paddle.CPUPlace()]
+        if not os.name == 'nt' and paddle.is_compiled_with_cuda():
+            self.places.append(paddle.CUDAPlace(0))
         paddle.enable_static()
 
     def tearDown(self):
         paddle.disable_static()
 
     def test_seed_without_tensor(self):
-        def api_run(seed):
+        def api_run(seed, place=paddle.CPUPlace()):
             main_prog, startup_prog = (
                 paddle.static.Program(),
                 paddle.static.Program(),
@@ -103,30 +106,35 @@ def api_run(seed):
             with paddle.static.program_guard(main_prog, startup_prog):
                 x = paddle.static.data(name='x', shape=[-1, 4], dtype='float32')
                 out = paddle.incubate.layers.shuffle_batch(x, seed=seed)
-            exe = paddle.static.Executor()
+            exe = paddle.static.Executor(place=place)
             feed = {'x': np.random.random((10, 4)).astype('float32')}
             exe.run(startup_prog)
             _ = exe.run(main_prog, feed=feed, fetch_list=[out])
 
-        api_run(seed=None)
-        api_run(seed=1)
+        for place in self.places:
+            api_run(None, place=place)
+            api_run(1, place=place)
 
     def test_seed_with_tensor(self):
-        main_prog, startup_prog = (
-            paddle.static.Program(),
-            paddle.static.Program(),
-        )
-        with paddle.static.program_guard(main_prog, startup_prog):
-            x = paddle.static.data(name='x', shape=[-1, 4], dtype='float32')
-            seed = paddle.static.data(name='seed', shape=[1], dtype='int64')
-            out = paddle.incubate.layers.shuffle_batch(x, seed=seed)
-        exe = paddle.static.Executor()
-        feed = {
-            'x': np.random.random((10, 4)).astype('float32'),
-            'seed': np.array([1]).astype('int64'),
-        }
-        exe.run(startup_prog)
-        _ = exe.run(main_prog, feed=feed, fetch_list=[out])
+        def api_run(place=paddle.CPUPlace()):
+            main_prog, startup_prog = (
+                paddle.static.Program(),
+                paddle.static.Program(),
+            )
+            with paddle.static.program_guard(main_prog, startup_prog):
+                x = paddle.static.data(name='x', shape=[-1, 4], dtype='float32')
+                seed = paddle.static.data(name='seed', shape=[1], dtype='int64')
+                out = paddle.incubate.layers.shuffle_batch(x, seed=seed)
+            exe = paddle.static.Executor(place=place)
+            feed = {
+                'x': np.random.random((10, 4)).astype('float32'),
+                'seed': np.array([1]).astype('int64'),
+            }
+            exe.run(startup_prog)
+            _ = exe.run(main_prog, feed=feed, fetch_list=[out])
+
+        for place in self.places:
+            api_run(place=place)
 
 
 if __name__ == '__main__':