diff options
| author | Boris Fomitchev <bfomitchev@nvidia.com> | 2018-05-08 20:18:10 -0700 |
|---|---|---|
| committer | Boris Fomitchev <bfomitchev@nvidia.com> | 2018-05-08 20:18:10 -0700 |
| commit | 25e205604e7eafa83867a15cfda526461fe58455 (patch) | |
| tree | fcce10851fb0d1627b60cc23100659506f1462bb /models/pix2pixHD_model.py | |
| parent | 4ca6b1610f9fa65f8bd7d7c15059bfde18a2f02a (diff) | |
ONNX export is working
Diffstat (limited to 'models/pix2pixHD_model.py')
| -rwxr-xr-x | models/pix2pixHD_model.py | 7 |
1 files changed, 7 insertions, 0 deletions
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) + + |