mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow tecogan to support generators that only output a tensor (instead of a list)

Browse Source
This commit is contained in:
James Betker 2020-10-08 09:26:25 -06:00
parent 969bcd9021
commit c174ac0fd5
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
codes/models/steps/tecogan_losses.py
View File

@ -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)