1 files changed, 2 insertions, 6 deletions
diff --git a/modules/lowvram.py b/modules/lowvram.py
index bd117491..079386c3 100644
--- a/modules/lowvram.py
+++ b/modules/lowvram.py
@@ -1,13 +1,9 @@
 import torch
+from modules.devices import get_optimal_device
 
 module_in_gpu = None
 cpu = torch.device("cpu")
-if torch.has_cuda:
-    device = gpu = torch.device("cuda")
-elif torch.has_mps:
-    device = gpu = torch.device("mps")
-else:
-    device = gpu = torch.device("cpu")
+device = gpu = get_optimal_device()
 
 def setup_for_low_vram(sd_model, use_medvram):
     parents = {}