fixing ssim metric

pptrick · 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