Переглянути джерело

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())