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 years ago
parent
commit
5960eb7539
1 changed files with 1 additions and 1 deletions
  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())