Browse Source

Fix the TypeError: expected CUDA (got CPU) error

The error is because the model is on GPU while your input image is on CPU. We have to make sure that they are both on GPU or CPU.
Student B 5 năm trước cách đây
mục cha
commit
5960eb7539
1 tập tin đã thay đổi với 1 bổ sung1 xóa
  1. 1 1
      src/rembg/u2net/detect.py

+ 1 - 1
src/rembg/u2net/detect.py

@@ -107,7 +107,7 @@ def predict(net, item):
     with torch.no_grad():
     with torch.no_grad():
 
 
         if torch.cuda.is_available():
         if torch.cuda.is_available():
-            inputs_test = torch.cuda.FloatTensor(sample["image"].unsqueeze(0).float())
+            inputs_test = torch.cuda.FloatTensor(sample["image"].unsqueeze(0).cuda().float())
         else:
         else:
             inputs_test = torch.FloatTensor(sample["image"].unsqueeze(0).float())
             inputs_test = torch.FloatTensor(sample["image"].unsqueeze(0).float())