From 25e205604e7eafa83867a15cfda526461fe58455 Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Tue, 8 May 2018 20:18:10 -0700 Subject: ONNX export is working --- models/models.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) (limited to 'models/models.py') diff --git a/models/models.py b/models/models.py index 805696f..8e72e46 100755 --- a/models/models.py +++ b/models/models.py @@ -4,8 +4,11 @@ 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() -- cgit v1.2.3-70-g09d2