Merge pull request AUTOMATIC1111#5095 from mlmcgoogan/master

torch.cuda.empty_cache() defaults to cuda:0 device unless explicitly …
HexPang · Nov 27, 2022 · 997ac57 · 997ac57
2 parents d860b56 + c67c40f
commit 997ac57
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/modules/devices.py b/modules/devices.py
@@ -44,8 +44,18 @@ def get_optimal_device():
 
 def torch_gc():
     if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
+        from modules import shared
+
+        device_id = shared.cmd_opts.device_id
+
+        if device_id is not None:
+            cuda_device = f"cuda:{device_id}"
+        else:
+            cuda_device = "cuda"
+
+        with torch.cuda.device(cuda_device):
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
 
 
 def enable_tf32():