fix half precision for cpu.

This commit is contained in:
sczhou committed 2023-11-09 22:09:24 +08:00
1 parent 3564cd3404
commit 9243b7abf1
2 files changed
+4 -1

No files matched your search

+2 -1
View File
@@ -218,7 +218,8 @@ if __name__ == '__main__':
# Use fp16 precision during inference to reduce running memory cost
use_half = True if args.fp16 else False
if device == torch.device('cpu'):
use_half = False
frames, fps, size, video_name = read_frame_from_videos(args.video)
if not args.width == -1 and not args.height == -1:
@@ -164,6 +164,8 @@ class ProInpainter:
def __init__(self, propainter_checkpoint, raft_checkpoint, flow_completion_checkpoint, device="cuda:0", use_half=True):
self.device = device
self.use_half = use_half
if self.device == torch.device('cpu'):
self.use_half = False
##############################################
# set up RAFT and flow competition model