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

Handle tuple-returning generators in test

Browse Source
This commit is contained in:
James Betker 2020-05-11 11:15:26 -06:00
parent f994466289
commit 62a97c53d1
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
codes/test.py
View File

@ -55,7 +55,10 @@ if __name__ == "__main__":
model.feed_data(data, need_GT=need_GT) model.feed_data(data, need_GT=need_GT)
model.test() model.test()
visuals = model.fake_H.detach().float().cpu() if isinstance(model.fake_H, tuple):
visuals = model.fake_H[0].detach().float().cpu()
else:
visuals = model.fake_H.detach().float().cpu()
for i in range(visuals.shape[0]): for i in range(visuals.shape[0]):
img_path = data['GT_path'][i] if need_GT else data['LQ_path'][i] img_path = data['GT_path'][i] if need_GT else data['LQ_path'][i]
img_name = osp.splitext(osp.basename(img_path))[0] img_name = osp.splitext(osp.basename(img_path))[0]