Merge pull request #4919 from brkirch/deepbooru-fix
Fix support for devices other than CUDA in DeepBooru
This commit is contained in:
commit
ef567b083c
|
@ -58,7 +58,7 @@ class DeepDanbooru:
|
||||||
a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
|
a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
|
||||||
|
|
||||||
with torch.no_grad(), devices.autocast():
|
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()
|
y = self.model(x)[0].detach().cpu().numpy()
|
||||||
|
|
||||||
probability_dict = {}
|
probability_dict = {}
|
||||||
|
|
Loading…
Reference in New Issue
Block a user