first commit

1 files changed, 14 insertions, 0 deletions

diff --git a/models/models.py b/models/models.py
new file mode 100755
index 0000000..351483c
--- /dev/null
+++ b/models/models.py
@@ -0,0 +1,14 @@
+### Copyright (C) 2017 NVIDIA Corporation. All rights reserved. 
+### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+import torch
+
+def create_model(opt):
+    from .pix2pixHD_model import Pix2PixHDModel
+    model = Pix2PixHDModel()    
+    model.initialize(opt)
+    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)
+
+    return model