From b60621f3d1ddace2588535050f9855ba92fde068 Mon Sep 17 00:00:00 2001
From: x54-729 <17307130121@fudan.edu.cn>
Date: Fri, 17 Jun 2022 23:23:43 +0800
Subject: [PATCH] small

---
 .../paddle_driver/test_initialize_paddle_driver.py       | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tests/core/drivers/paddle_driver/test_initialize_paddle_driver.py b/tests/core/drivers/paddle_driver/test_initialize_paddle_driver.py
index 7e567c84..63124cdc 100644
--- a/tests/core/drivers/paddle_driver/test_initialize_paddle_driver.py
+++ b/tests/core/drivers/paddle_driver/test_initialize_paddle_driver.py
@@ -1,3 +1,5 @@
+import os
+
 import pytest
 
 from fastNLP.core.drivers import PaddleSingleDriver, PaddleFleetDriver
@@ -40,9 +42,14 @@ def test_get_fleet(device):
     """
     测试 fleet 多卡的初始化情况
     """
-
+    flag = False
+    if "USER_CUDA_VISIBLE_DEVICES" not in os.environ:
+        os.environ["USER_CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
+        flag = True
     model = PaddleNormalModel_Classification_1(20, 10)
     driver = initialize_paddle_driver("paddle", device, model)
+    if flag:
+        del os.environ["USER_CUDA_VISIBLE_DEVICES"]
 
     assert isinstance(driver, PaddleFleetDriver)