Skip to content

Commit 208d4f6

Browse files
committed
Modify cuda argument type
1 parent 22c0efb commit 208d4f6

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

demo.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
parser.add_argument('--style_image_path', default='./images/style1.png')
2020
parser.add_argument('--style_seg_path', default=[])
2121
parser.add_argument('--output_image_path', default='./results/example1.png')
22-
parser.add_argument('--cuda', type=bool, default=True, help='Enable CUDA.')
22+
parser.add_argument('--cuda', type=int, default=1, help='Enable CUDA.')
2323
args = parser.parse_args()
2424

2525
# Load model
@@ -30,7 +30,8 @@
3030
print("Fail to load PhotoWCT models. PhotoWCT submodule not updated?")
3131
exit()
3232

33-
p_wct.cuda(0)
33+
if args.cuda:
34+
p_wct.cuda(0)
3435

3536
process_stylization.stylization(
3637
p_wct=p_wct,

0 commit comments

Comments
 (0)