fix SSIM averaging bug

miaotianyi · Jun 24, 2021 · 5434f41 · 5434f41
1 parent 23a2a15
commit 5434f41
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchimage/metrics/ssim_new.py b/torchimage/metrics/ssim_new.py
@@ -83,7 +83,7 @@ def _crop_border(self, full: torch.Tensor, axes: tuple):
         # access last-used padding width
         for i, a in enumerate(axes):
             pad_before, pad_after = pad_width[i]
-            idx[i] = slice(pad_before, full.shape[a]-pad_after)
+            idx[a] = slice(pad_before, full.shape[a]-pad_after)
         return full[tuple(idx)]
 
     def forward_score(self, x: torch.Tensor, content_axes: tuple, avg_axes: tuple):