fixing ssim metric

stalhabukhari · Feb 15, 2023 · e659839 · e659839
1 parent f4c3486
commit e659839
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nerf/utils.py b/nerf/utils.py
@@ -241,10 +241,12 @@ def report(self):
 
 
 class SSIMMeter:
-    def __init__(self):
+    def __init__(self, device=None):
         self.V = 0
         self.N = 0
 
+        self.device = device if device is not None else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
     def clear(self):
         self.V = 0
         self.N = 0