From ab99580bd90fa03dfa2df163bd40b6ca9226605d Mon Sep 17 00:00:00 2001
From: hzwer <598460606@163.com>
Date: Sun, 15 Nov 2020 13:18:09 +0800
Subject: [PATCH] Use CUDA to evaluation

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

diff --git a/Vimeo90K_benchmark.py b/Vimeo90K_benchmark.py
index 2b00706..d5c9d54 100644
--- a/Vimeo90K_benchmark.py
+++ b/Vimeo90K_benchmark.py
@@ -33,7 +33,7 @@ for i in f:
     mid = np.round((mid * 255).cpu().numpy()).astype('uint8').transpose(1, 2, 0) / 255.
     I1 = I1 / 255.
     psnr = -10 * math.log10(((I1 - mid) * (I1 - mid)).mean())
-    ssim = ssim_matlab(torch.tensor(I1).unsqueeze(0).float(), torch.tensor(mid).unsqueeze(0).float())
+    ssim = ssim_matlab(torch.tensor(I1).unsqueeze(0).float().to(device), torch.tensor(mid).unsqueeze(0).float().to(device)).cpu().numpy()
     psnr_list.append(psnr)
     ssim_list.append(ssim)
     print(np.mean(psnr_list), np.mean(ssim_list))