From c174ac0fd55b1880cc4354ac07f1a264e523788b Mon Sep 17 00:00:00 2001 From: James Betker Date: Thu, 8 Oct 2020 09:26:25 -0600 Subject: [PATCH] Allow tecogan to support generators that only output a tensor (instead of a list) --- codes/models/steps/tecogan_losses.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/codes/models/steps/tecogan_losses.py b/codes/models/steps/tecogan_losses.py index 5857a5f5..ea8b1017 100644 --- a/codes/models/steps/tecogan_losses.py +++ b/codes/models/steps/tecogan_losses.py @@ -95,6 +95,8 @@ class RecurrentImageGeneratorSequenceInjector(Injector): self.produce_teco_visual_debugs(input[self.input_lq_index], debug_index) debug_index += 1 gen_out = gen(*input) + if isinstance(gen_out, torch.Tensor): + gen_out = [gen_out] recurrent_input = gen_out[self.output_hq_index] results.append(recurrent_input) @@ -113,6 +115,8 @@ class RecurrentImageGeneratorSequenceInjector(Injector): self.produce_teco_visual_debugs(input[self.input_lq_index], debug_index) debug_index += 1 gen_out = gen(*input) + if isinstance(gen_out, torch.Tensor): + gen_out = [gen_out] recurrent_input = gen_out[self.output_hq_index] results.append(recurrent_input)