From 25444da59d0838ffd30453a8b9c8d62dbe1e695e Mon Sep 17 00:00:00 2001 From: Somdev Sangwan Date: Tue, 30 May 2023 07:00:37 +0530 Subject: [PATCH] fix the --gpu mess --- run.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/run.py b/run.py index 1cecab2..aceab00 100755 --- a/run.py +++ b/run.py @@ -36,13 +36,14 @@ parser.add_argument('-f', '--face', help='use this face', dest='source_img') parser.add_argument('-t', '--target', help='replace this face', dest='target_path') parser.add_argument('-o', '--output', help='save output to this file', dest='output_file') parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False) -if torch.cuda.is_available(): - parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False) +parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) for name, value in vars(parser.parse_args()).items(): args[name] = value +if torch.cuda.is_available() not args['gpu']: + quit("You are using --gpu flag but CUDA isn't available on your system.") sep = "/" if os.name == "nt":