From f194457229e4537912467bc60ac3a873f473a63c Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 11 Sep 2022 18:48:36 +0300
Subject: CLIP interrogator

---
 modules/devices.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

(limited to 'modules/devices.py')

diff --git a/modules/devices.py b/modules/devices.py
index 25008a04..30d30b99 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -1,12 +1,16 @@
 import torch
 
-
 # has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
 has_mps = getattr(torch, 'has_mps', False)
 
+cpu = torch.device("cpu")
+
+
 def get_optimal_device():
-  if torch.cuda.is_available():
-      return torch.device("cuda")
-  if has_mps:
-      return torch.device("mps")
-  return torch.device("cpu")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+
+    if has_mps:
+        return torch.device("mps")
+
+    return cpu
-- 
cgit v1.2.3