From 421588e9250b3fe5e8c4ad6f6c2bcb95fac28aeb Mon Sep 17 00:00:00 2001
From: hzwer <598460606@163.com>
Date: Mon, 16 Nov 2020 17:55:21 +0800
Subject: [PATCH] Fix inference

---
 inference_img.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/inference_img.py b/inference_img.py
index 9ca6fac..d9b78cb 100644
--- a/inference_img.py
+++ b/inference_img.py
@@ -42,4 +42,4 @@ for i in range(args.times):
 if not os.path.exists('output'):
     os.mkdir('output')
 for i in range(len(img_list)):
-    cv2.imwrite('output/img{}.png'.format(i), img_list[i][0].numpy().transpose(1, 2, 0)[:h, :w] * 255)
+    cv2.imwrite('output/img{}.png'.format(i), img_list[i][0].cpu().numpy().transpose(1, 2, 0)[:h, :w] * 255)