From 563ea3f6ff66e0eba44033163d24e42297465a47 Mon Sep 17 00:00:00 2001 From: brkirch Date: Mon, 21 Nov 2022 02:56:00 -0500 Subject: [PATCH] Change .cuda() to .to(devices.device) --- modules/deepbooru.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/deepbooru.py b/modules/deepbooru.py index b9066d817..31ec7e171 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 = {}