Skip to content

Commit 563ea3f

Browse files
committed
Change .cuda() to .to(devices.device)
1 parent 828438b commit 563ea3f

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

modules/deepbooru.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ def tag_multi(self, pil_image, force_disable_ranks=False):
5858
a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
5959

6060
with torch.no_grad(), devices.autocast():
61-
x = torch.from_numpy(a).cuda()
61+
x = torch.from_numpy(a).to(devices.device)
6262
y = self.model(x)[0].detach().cpu().numpy()
6363

6464
probability_dict = {}

0 commit comments

Comments
 (0)