From 446322754aaa6b6ccaabbc1d71b7993554252c58 Mon Sep 17 00:00:00 2001 From: James Betker Date: Sat, 23 May 2020 21:09:54 -0600 Subject: [PATCH] Support generators that don't output intermediary values. --- codes/models/SRGAN_model.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/codes/models/SRGAN_model.py b/codes/models/SRGAN_model.py index ba0f3ade..66fc05a6 100644 --- a/codes/models/SRGAN_model.py +++ b/codes/models/SRGAN_model.py @@ -327,10 +327,11 @@ class SRGANModel(BaseModel): utils.save_image(self.pix[i].cpu().detach(), os.path.join(sample_save_path, "pix", "%05i_%02i.png" % (step, i))) if multi_gen: utils.save_image(self.fake_GenOut[i][0].cpu().detach(), os.path.join(sample_save_path, "gen", "%05i_%02i.png" % (step, i))) - if self.fake_GenOut[i][1] is not None: - utils.save_image(self.fake_GenOut[i][1].cpu().detach(), os.path.join(sample_save_path, "genmr", "%05i_%02i.png" % (step, i))) - if self.fake_GenOut[i][2] is not None: - utils.save_image(self.fake_GenOut[i][2].cpu().detach(), os.path.join(sample_save_path, "genlr", "%05i_%02i.png" % (step, i))) + if len(self.fake_GenOut[i]) > 1: + if self.fake_GenOut[i][1] is not None: + utils.save_image(self.fake_GenOut[i][1].cpu().detach(), os.path.join(sample_save_path, "genmr", "%05i_%02i.png" % (step, i))) + if self.fake_GenOut[i][2] is not None: + utils.save_image(self.fake_GenOut[i][2].cpu().detach(), os.path.join(sample_save_path, "genlr", "%05i_%02i.png" % (step, i))) utils.save_image(var_ref_skips[i][0].cpu().detach(), os.path.join(sample_save_path, "ref", "hi_%05i_%02i.png" % (step, i))) utils.save_image(var_ref_skips[i][1].cpu().detach(), os.path.join(sample_save_path, "ref", "med_%05i_%02i.png" % (step, i))) utils.save_image(var_ref_skips[i][2].cpu().detach(), os.path.join(sample_save_path, "ref", "low_%05i_%02i.png" % (step, i)))