Merge pull request #4919 from brkirch/deepbooru-fix

Fix support for devices other than CUDA in DeepBooru
This commit is contained in:
AUTOMATIC1111 2022-11-27 16:59:22 +03:00 committed by GitHub
commit ef567b083c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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 = {}