torch check nvidia only

This commit is contained in:
Phan Tuấn Anh 2023-06-05 12:25:35 +02:00 committed by GitHub
parent 3668360520
commit b947603441
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -209,7 +209,7 @@ def start(preview_callback = None):
process_video_multi_cores(args.source_img, args.frame_paths) process_video_multi_cores(args.source_img, args.frame_paths)
else: else:
process_video(args.source_img, args.frame_paths) process_video(args.source_img, args.frame_paths)
if torch.cuda.is_available() and args.gpu_vendor in ['amd', 'nvidia']: if args.gpu_vendor = 'nvidia':
torch.cuda.empty_cache() torch.cuda.empty_cache()
status("creating video...") status("creating video...")
create_video(video_name, exact_fps, output_dir) create_video(video_name, exact_fps, output_dir)