Prioritize Previously-Used GPUs During Auto-Selection (#10485)

Co-authored-by: Rohit Gupta <[email protected]>
Co-authored-by: Adrian Wälchli <[email protected]>

Author: 3 people

pytorch_lightning/tuner/auto_gpu_select.py

@@ -45,9 +45,19 @@ def pick_single_gpu(exclude_gpus: List[int]) -> int:
         RuntimeError:
             If you try to allocate a GPU, when no GPUs are available.
     """
+    previously_used_gpus = []
+    unused_gpus = []
     for i in range(torch.cuda.device_count()):
         if i in exclude_gpus:
             continue
+
+        if torch.cuda.memory_reserved(f"cuda:{i}") > 0:
+            previously_used_gpus.append(i)
+        else:
+            unused_gpus.append(i)
+
+    # Prioritize previously used GPUs
+    for i in previously_used_gpus + unused_gpus:
         # Try to allocate on device:
         device = torch.device(f"cuda:{i}")
         try: