N2V2 fix

src/careamics/models/unet.py

@@ -180,11 +180,7 @@ def __init__(
         decoder_blocks = []
         for n in range(depth):
             decoder_blocks.append(upsampling)
-            in_channels = (
-                num_channels_init ** (depth - n)
-                if (self.n2v2 and n == depth - 1)
-                else num_channels_init * 2 ** (depth - n)
-            )
+            in_channels = num_channels_init * 2 ** (depth - n)
             out_channels = in_channels // 2
             decoder_blocks.append(
                 Conv_Block(

tests/test_lightning_module.py

@@ -167,6 +167,40 @@ def test_careamics_kiln_unet_depth_3_3D(shape):
     assert y.shape == x.shape
 
 
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (8, 64, 64),
+        (16, 64, 64),
+        (16, 128, 128),
+        (32, 128, 128),
+    ],
+)
+def test_careamics_kiln_unet_depth_3_3D_n2v2(shape):
+    algo_dict = {
+        "algorithm": "n2v",
+        "model": {
+            "architecture": "UNet",
+            "conv_dims": 3,
+            "in_channels": 1,
+            "num_classes": 1,
+            "depth": 3,
+            "n2v2": True,
+        },
+        "loss": "n2v",
+    }
+    algo_config = AlgorithmConfig(**algo_dict)
+
+    # instantiate CAREamicsKiln
+    model = CAREamicsModule(algo_config)
+    # set model to evaluation mode to avoid batch dimension error
+    model.model.eval()
+    # test forward pass
+    x = torch.rand((1, 1, *shape))
+    y: torch.Tensor = model.forward(x)
+    assert y.shape == x.shape
+
+
 @pytest.mark.parametrize("n_channels", [1, 3, 4])
 def test_careamics_kiln_unet_depth_2_channels_2D(n_channels):
     algo_dict = {