From a3355ca65995ba244a6ae187ca66d4841b06aa2c Mon Sep 17 00:00:00 2001
From: OodavidsinoO <david.sin@tuta.io>
Date: Thu, 5 May 2022 10:24:16 +0800
Subject: [PATCH] Update torch.device() logic

---
 Disco_Diffusion.ipynb | 2 +-
 disco.py              | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/Disco_Diffusion.ipynb b/Disco_Diffusion.ipynb
index 17493be..b5b1f33 100644
--- a/Disco_Diffusion.ipynb
+++ b/Disco_Diffusion.ipynb
@@ -596,7 +596,7 @@
         "  MAX_ADABINS_AREA = 500000\n",
         "\n",
         "import torch\n",
-        "DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
+        "DEVICE = torch.device('cuda:0' if (torch.cuda.is_available() and not useCPU) else 'cpu')\n",
         "print('Using device:', DEVICE)\n",
         "device = DEVICE # At least one of the modules expects this name..\n",
         "\n",
diff --git a/disco.py b/disco.py
index 254c57a..11711c1 100644
--- a/disco.py
+++ b/disco.py
@@ -569,7 +569,7 @@ if USE_ADABINS:
   MAX_ADABINS_AREA = 500000
 
 import torch
-DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+DEVICE = torch.device('cuda:0' if (torch.cuda.is_available() and not useCPU) else 'cpu')
 print('Using device:', DEVICE)
 device = DEVICE # At least one of the modules expects this name..