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/pix2pixHD_model.py | 7 +++++++ 1 file changed, 7 insertions(+) (limited to 'models/pix2pixHD_model.py') 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) + + -- cgit v1.2.3-70-g09d2