fix: use torch to detect GPU number instead of pynvml

RUC-NLPIR · Dec 10, 2024 · 6a8d968 · 6a8d968
1 parent 3e8c0ad
commit 6a8d968
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/flashrag/config/config.py b/flashrag/config/config.py
@@ -104,9 +104,11 @@ def _init_device(self):
         if gpu_id is not None:
             os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
         try:
-            import pynvml 
-            pynvml.nvmlInit()
-            gpu_num = pynvml.nvmlDeviceGetCount()
+            # import pynvml 
+            # pynvml.nvmlInit()
+            # gpu_num = pynvml.nvmlDeviceGetCount()
+            import torch
+            gpu_num = torch.cuda.device_count()
         except:
             gpu_num = 0
         self.final_config['gpu_num'] = gpu_num