From 9054cf9b0c327a5077fd0793abe178f400da3315 Mon Sep 17 00:00:00 2001 From: tingchunw Date: Mon, 4 Dec 2017 16:52:46 -0800 Subject: first commit --- models/models.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100755 models/models.py (limited to 'models/models.py') 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 -- cgit v1.2.3-70-g09d2