From edf910b1c1d02020b31782ab4c3b6ebf9af8c323 Mon Sep 17 00:00:00 2001
From: tingchunw <tingchunw@nvidia.com>
Date: Tue, 30 Jan 2018 01:30:18 +0000
Subject: change dataset naming convention and add ui model

---
 models/models.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

(limited to 'models/models.py')

diff --git a/models/models.py b/models/models.py
index 351483c..0ba442f 100755
--- a/models/models.py
+++ b/models/models.py
@@ -3,8 +3,12 @@
 import torch
 
 def create_model(opt):
-    from .pix2pixHD_model import Pix2PixHDModel
-    model = Pix2PixHDModel()    
+    if opt.model == 'pix2pixHD':
+    	from .pix2pixHD_model import Pix2PixHDModel
+    	model = Pix2PixHDModel()    
+    else:
+    	from .ui_model import UIModel
+    	model = UIModel()
     model.initialize(opt)
     print("model [%s] was created" % (model.name()))
 
-- 
cgit v1.2.3-70-g09d2