ONNX export is working

2 files changed, 12 insertions, 2 deletions

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()
diff --git a/models/pix2pixHD_model.py b/models/pix2pixHD_model.py
index de594ab..631a10f 100755
--- a/models/pix2pixHD_model.py
+++ b/models/pix2pixHD_model.py
@@ -270,3 +270,10 @@ class Pix2PixHDModel(BaseModel):
         if self.opt.verbose:
             print('update learning rate: %f -> %f' % (self.old_lr, lr))
         self.old_lr = lr
+
+class InferenceModel(Pix2PixHDModel):
+    def forward(self, inp):
+        label, inst = inp
+        return self.inference(label, inst)
+
+