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..