修改测试中需要超过2个gpu的

fastNLP/core/callbacks/topk_saver.py

@@ -36,7 +36,7 @@ class Saver:
                  model_save_fn:Callable=None, **kwargs):
         if folder is None:
             folder = Path.cwd().absolute()
-            logger.info(f"Parameter `folder` is None, and we will use {folder} to save and load your model.")
+            logger.info(f"Parameter `folder` is None, and fastNLP will use {folder} to save and load your model.")
         folder = Path(folder)
         if not folder.exists():
             folder.mkdir(parents=True, exist_ok=True)

tests/core/controllers/_test_distributed_launch_torch_1.py

@@ -6,7 +6,7 @@ python -m torch.distributed.launch --nproc_per_node 2 tests/core/controllers/_te
 
 import argparse
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1"
 
 import sys
 path = os.path.abspath(__file__)

tests/core/controllers/test_trainer_event_trigger.py

@@ -224,7 +224,7 @@ def test_trainer_event_trigger_2(
                 assert k in output[0]
 
 
-@pytest.mark.parametrize("driver,device", [("torch", "cpu"), ("torch", 6)])
+@pytest.mark.parametrize("driver,device", [("torch", "cpu"), ("torch", 0)])
 @pytest.mark.torch
 @magic_argv_env_context
 def test_trainer_event_trigger_3(

tests/core/dataloaders/jittor_dataloader/test_fdl.py

@@ -1,6 +1,9 @@
 import pytest
 import numpy as np
-from datasets import Dataset as HfDataset
+from fastNLP.envs import _module_available
+
+if _module_available('datasets'):
+    from datasets import Dataset as HfDataset
 
 from fastNLP.core.dataloaders.jittor_dataloader import JittorDataLoader
 from fastNLP.core.dataset import DataSet as Fdataset

tests/core/drivers/torch_driver/test_initialize_torch_driver.py

@@ -40,7 +40,7 @@ def test_get_single_device(driver, device):
 @pytest.mark.torch
 @pytest.mark.parametrize(
     "device", 
-    [[0, 2, 3], -1]
+    [[0, 1], -1]
 )
 @pytest.mark.parametrize(
     "driver", 

tests/core/metrics/test_accuracy_torch.py

@@ -102,7 +102,7 @@ class TestAccuracy:
                     metric_kwargs=metric_kwargs,
                     sklearn_metric=sklearn_accuracy,
                 ),
-                [(rank, processes, torch.device(f'cuda:{rank+4}')) for rank in range(processes)]
+                [(rank, processes, torch.device(f'cuda:{rank}')) for rank in range(processes)]
             )
         else:
             device = torch.device(

tests/core/metrics/test_classify_f1_pre_rec_metric_torch.py

@@ -177,6 +177,6 @@ class TestClassfiyFPreRecMetric:
                              metric_class=ClassifyFPreRecMetric,
                              metric_kwargs=metric_kwargs,
                              metric_result=ground_truth),
-                     [(rank, NUM_PROCESSES, torch.device(f'cuda:{rank+4}')) for rank in range(NUM_PROCESSES)])
+                     [(rank, NUM_PROCESSES, torch.device(f'cuda:{rank}')) for rank in range(NUM_PROCESSES)])
         pool.close()
         pool.join()