Repository: Project-MONAI/MONAI

Issue:
SSIM loss fail when batchsize>1
**Describe the bug**
SSIM loss fail when batchsize>1

**To Reproduce**
Steps to reproduce the behavior:
```
import torch
from monai.losses.ssim_loss import SSIMLoss
x = torch.ones([4,1,10,10])/2
y = torch.ones([4,1,10,10])/2
data_range = x.max().unsqueeze(0)
print(1-SSIMLoss(spatial_dims=2)(x,y,data_range))
```

**Expected behavior**
Should print 1, but got error

**Screenshots**
If applicable, add screenshots to help explain your problem.

**Environment**

Ensuring you use the relevant python executable, please paste the output of:

```
python -c 'import monai; monai.config.print_debug_info()'
```

**Additional context**
Add any other context about the problem here.



Discussion:
Hi @PedroFerreiradaCosta, I was wondering if you would have time to look at it. If not, I could work on it in two weeks. Thank you a lot!

Patch:
diff --git a/monai/losses/ssim_loss.py b/monai/losses/ssim_loss.py
--- a/monai/losses/ssim_loss.py
+++ b/monai/losses/ssim_loss.py
@@ -95,7 +95,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor, data_range: torch.Tensor) ->
                 if i == 0:
                     ssim_value = ssim_val
                 else:
-                    ssim_value = torch.cat((ssim_value.view(1), ssim_val.view(1)), dim=0)
+                    ssim_value = torch.cat((ssim_value.view(i), ssim_val.view(1)), dim=0)
 
         else:
             raise ValueError("Batch size is not nonnegative integer value")


Test patch:
diff --git a/tests/test_ssim_loss.py b/tests/test_ssim_loss.py
--- a/tests/test_ssim_loss.py
+++ b/tests/test_ssim_loss.py
@@ -18,31 +18,34 @@
 
 from monai.losses.ssim_loss import SSIMLoss
 
-x = torch.ones([1, 1, 10, 10]) / 2
-y1 = torch.ones([1, 1, 10, 10]) / 2
-y2 = torch.zeros([1, 1, 10, 10])
-data_range = x.max().unsqueeze(0)
 TESTS2D = []
 for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]:
-    TESTS2D.append((x.to(device), y1.to(device), data_range.to(device), torch.tensor(1.0).unsqueeze(0).to(device)))
-    TESTS2D.append((x.to(device), y2.to(device), data_range.to(device), torch.tensor(0.0).unsqueeze(0).to(device)))
+    for batch_size in [1, 2, 16]:
+        x = torch.ones([batch_size, 1, 10, 10]) / 2
+        y1 = torch.ones([batch_size, 1, 10, 10]) / 2
+        y2 = torch.zeros([batch_size, 1, 10, 10])
+        data_range = x.max().unsqueeze(0)
+        TESTS2D.append((x.to(device), y1.to(device), data_range.to(device), torch.tensor(1.0).unsqueeze(0).to(device)))
+        TESTS2D.append((x.to(device), y2.to(device), data_range.to(device), torch.tensor(0.0).unsqueeze(0).to(device)))
 
-x = torch.ones([1, 1, 10, 10, 10]) / 2
-y1 = torch.ones([1, 1, 10, 10, 10]) / 2
-y2 = torch.zeros([1, 1, 10, 10, 10])
-data_range = x.max().unsqueeze(0)
 TESTS3D = []
 for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]:
-    TESTS3D.append((x.to(device), y1.to(device), data_range.to(device), torch.tensor(1.0).unsqueeze(0).to(device)))
-    TESTS3D.append((x.to(device), y2.to(device), data_range.to(device), torch.tensor(0.0).unsqueeze(0).to(device)))
+    for batch_size in [1, 2, 16]:
+        x = torch.ones([batch_size, 1, 10, 10, 10]) / 2
+        y1 = torch.ones([batch_size, 1, 10, 10, 10]) / 2
+        y2 = torch.zeros([batch_size, 1, 10, 10, 10])
+        data_range = x.max().unsqueeze(0)
+        TESTS3D.append((x.to(device), y1.to(device), data_range.to(device), torch.tensor(1.0).unsqueeze(0).to(device)))
+        TESTS3D.append((x.to(device), y2.to(device), data_range.to(device), torch.tensor(0.0).unsqueeze(0).to(device)))
 
-x = torch.ones([1, 1, 10, 10]) / 2
-y = torch.ones([1, 1, 10, 10]) / 2
-y.requires_grad_(True)
-data_range = x.max().unsqueeze(0)
 TESTS2D_GRAD = []
 for device in [None, "cpu", "cuda"] if torch.cuda.is_available() else [None, "cpu"]:
-    TESTS2D_GRAD.append([x.to(device), y.to(device), data_range.to(device)])
+    for batch_size in [1, 2, 16]:
+        x = torch.ones([batch_size, 1, 10, 10]) / 2
+        y = torch.ones([batch_size, 1, 10, 10]) / 2
+        y.requires_grad_(True)
+        data_range = x.max().unsqueeze(0)
+        TESTS2D_GRAD.append([x.to(device), y.to(device), data_range.to(device)])
 
 
 class TestSSIMLoss(unittest.TestCase):