fix the model forward bug caused by old PyTorch version

pocketdigi · Dec 17, 2020 · fea4e3a · fea4e3a
1 parent 60760ce
commit fea4e3a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/demo/image_matting/colab/inference.py b/demo/image_matting/colab/inference.py
@@ -90,7 +90,7 @@
         im = F.interpolate(im, size=(im_rh, im_rw), mode='area')
 
         # inference
-        _, _, matte = modnet(im.cuda(), inference=False)
+        _, _, matte = modnet(im.cuda(), True)
 
         # resize and save matte
         matte = F.interpolate(matte, size=(im_h, im_w), mode='area')

diff --git a/demo/video_matting/webcam/run.py b/demo/video_matting/webcam/run.py
@@ -52,7 +52,7 @@
         frame_tensor = frame_tensor.cuda()
 
     with torch.no_grad():
-        _, _, matte_tensor = modnet(frame_tensor, inference=True)
+        _, _, matte_tensor = modnet(frame_tensor, True)
 
     matte_tensor = matte_tensor.repeat(1, 3, 1, 1)
     matte_np = matte_tensor[0].data.cpu().numpy().transpose(1, 2, 0)