瀏覽代碼

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 年之前
父節點
當前提交
5960eb7539
共有 1 個文件被更改,包括 1 次插入1 次删除
  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():
 
         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:
             inputs_test = torch.FloatTensor(sample["image"].unsqueeze(0).float())