Merge pull request #4919 from brkirch/deepbooru-fix

Fix support for devices other than CUDA in DeepBooru

1 files changed, 1 insertions, 1 deletions

diff --git a/modules/deepbooru.py b/modules/deepbooru.py
index b9066d81..31ec7e17 100644
--- a/modules/deepbooru.py
+++ b/modules/deepbooru.py
@@ -58,7 +58,7 @@ class DeepDanbooru:
         a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
 
         with torch.no_grad(), devices.autocast():
-            x = torch.from_numpy(a).cuda()
+            x = torch.from_numpy(a).to(devices.device)
             y = self.model(x)[0].detach().cpu().numpy()
 
         probability_dict = {}