From 6dad10574b7d1caa6b3fd7e15d089dd8bc43d046 Mon Sep 17 00:00:00 2001
From: Wenjie Du <wenjay.du@gmail.com>
Date: Wed, 6 Dec 2023 21:37:26 +0800
Subject: [PATCH] refactor: using only one GPU during testing if there are more
 than one available;

---
 tests/global_test_config.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/global_test_config.py b/tests/global_test_config.py
index 3258ab6d..6d307427 100644
--- a/tests/global_test_config.py
+++ b/tests/global_test_config.py
@@ -7,6 +7,7 @@
 
 import os
 
+import numpy as np
 import torch
 
 from pypots.data.generating import gene_random_walk
@@ -33,10 +34,11 @@
 
 
 # set DEVICES to None if no cuda device is available, to avoid initialization failed while importing test classes
-cuda_devices = [torch.device(i) for i in range(torch.cuda.device_count())]
-if len(cuda_devices) > 2:
+n_cuda_devices = torch.cuda.device_count()
+cuda_devices = [torch.device(i) for i in range(n_cuda_devices)]
+if n_cuda_devices > 1:
     logger.info("❗️Detected multiple cuda devices, using all of them to run testing.")
-    DEVICE = cuda_devices
+    DEVICE = cuda_devices[np.random.randint(n_cuda_devices)]
 else:
     # if having no multiple cuda devices, leave it as None to use the default device
     DEVICE = None