From 208d4f60d036dd6944254e41e804f79806131c2b Mon Sep 17 00:00:00 2001
From: Ming-Yu Liu <mingyu.liu.tw@gmail.com>
Date: Fri, 2 Mar 2018 10:13:05 -0800
Subject: [PATCH] Modify cuda argument type

---
 demo.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/demo.py b/demo.py
index 8297f02..c1bb3bf 100644
--- a/demo.py
+++ b/demo.py
@@ -19,7 +19,7 @@
 parser.add_argument('--style_image_path', default='./images/style1.png')
 parser.add_argument('--style_seg_path', default=[])
 parser.add_argument('--output_image_path', default='./results/example1.png')
-parser.add_argument('--cuda', type=bool, default=True, help='Enable CUDA.')
+parser.add_argument('--cuda', type=int, default=1, help='Enable CUDA.')
 args = parser.parse_args()
 
 # Load model
@@ -30,7 +30,8 @@
     print("Fail to load PhotoWCT models. PhotoWCT submodule not updated?")
     exit()
 
-p_wct.cuda(0)
+if args.cuda:
+    p_wct.cuda(0)
 
 process_stylization.stylization(
     p_wct=p_wct,