From 32621480292ddaa29d2ebcba7fc025ac2449e1ff Mon Sep 17 00:00:00 2001 From: Zhanghan Ke Date: Thu, 17 Dec 2020 13:15:40 +0800 Subject: [PATCH] fix WebCam demo CPU bug --- demo/video_matting/webcam/run.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/demo/video_matting/webcam/run.py b/demo/video_matting/webcam/run.py index 3aaa1a3..79a482d 100644 --- a/demo/video_matting/webcam/run.py +++ b/demo/video_matting/webcam/run.py @@ -47,8 +47,10 @@ while(True): frame_PIL = Image.fromarray(frame_np) frame_tensor = torch_transforms(frame_PIL) - frame_tensor = frame_tensor[None, :, :, :].cuda() - + frame_tensor = frame_tensor[None, :, :, :] + if GPU: + frame_tensor = frame_tensor.cuda() + with torch.no_grad(): _, _, matte_tensor = modnet(frame_tensor, inference=True)