Merge pull request #33 from borisfom/fp16

Added data size and ONNX export options, FP16 inference is working

1 files changed, 7 insertions, 3 deletions

diff --git a/models/models.py b/models/models.py
index 0ba442f..8e72e46 100755
--- a/models/models.py
+++ b/models/models.py
@@ -4,13 +4,17 @@ import torch
 
 def create_model(opt):
     if opt.model == 'pix2pixHD':
-    	from .pix2pixHD_model import Pix2PixHDModel
-    	model = Pix2PixHDModel()    
+    	from .pix2pixHD_model import Pix2PixHDModel, InferenceModel
+    	if opt.isTrain:
+            model = Pix2PixHDModel()
+        else:
+            model = InferenceModel()
     else:
     	from .ui_model import UIModel
     	model = UIModel()
     model.initialize(opt)
-    print("model [%s] was created" % (model.name()))
+    if opt.verbose:
+        print("model [%s] was created" % (model.name()))
 
     if opt.isTrain and len(opt.gpu_ids):
         model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)