Update README.md

README.md

@@ -26,6 +26,7 @@ import warnings
 warnings.filterwarnings("ignore")
 
 PATH_TO_IMAGE = 'https://images.unsplash.com/photo-1594568284297-7c64464062b1'
+USE_CUDA = False
 
 def download_model():
     print("Downloading google_safesearch_mini.bin...")
@@ -52,8 +53,12 @@ def run():
     ])
     img = transform(img)
     img = img.unsqueeze(0)
-    img = img.cpu()
-    model = model.cpu()
+    if USE_CUDA:
+        img = img.cuda()
+        model = model.cuda()
+    else:
+        img = img.cpu()
+        model = model.cpu()
     model.eval()
     with torch.no_grad():
         out, _ = model(img)