From df109316e5e0d1fae8e52df7da336ddfc1f1b508 Mon Sep 17 00:00:00 2001
From: x54-729 <17307130121@fudan.edu.cn>
Date: Tue, 26 Apr 2022 05:33:37 +0000
Subject: [PATCH] small

---
 tests/core/drivers/paddle_driver/test_fleet.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/core/drivers/paddle_driver/test_fleet.py b/tests/core/drivers/paddle_driver/test_fleet.py
index 52739f53..34c80888 100644
--- a/tests/core/drivers/paddle_driver/test_fleet.py
+++ b/tests/core/drivers/paddle_driver/test_fleet.py
@@ -527,7 +527,7 @@ class TestSaveLoad:
     @classmethod
     def setup_class(cls):
         # 不在这里 setup 的话会报错
-        cls.driver = generate_driver(10, 10)
+        cls.driver = generate_driver(10, 10, device=[0,1])
 
     def setup_method(self):
         self.dataset = PaddleRandomMaxDataset(20, 10)
@@ -633,7 +633,7 @@ class TestSaveLoad:
                 batch_sampler=BucketedBatchSampler(
                     self.dataset,
                     length=[10 for i in range(len(self.dataset))],
-                    batch_size=4,
+                    batch_size=2,
                 )
             )
             dataloader.batch_sampler.set_distributed(