From 62a97c53d115e05835880da6dbff1f0482a434ed Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Mon, 11 May 2020 11:15:26 -0600
Subject: [PATCH] Handle tuple-returning generators in test
---
codes/test.py | 5 ++++-
1 file changed, 4 insertions(+), 1 deletion(-)
diff --git a/codes/test.py b/codes/test.py
index 4f4498a3..a764d3f8 100644
--- a/codes/test.py
+++ b/codes/test.py
@@ -55,7 +55,10 @@ if __name__ == "__main__":
model.feed_data(data, need_GT=need_GT)
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]):
img_path = data['GT_path'][i] if need_GT else data['LQ_path'][i]
img_name = osp.splitext(osp.basename(img_path))[0]