add toggle for using skip connections

models/convnext.py

@@ -180,9 +180,12 @@ class ConvNeXtReverseBase2(nn.Module):
         block: Optional[Callable[..., nn.Module]] = None,
         norm_layer: Optional[Callable[..., nn.Module]] = None,
         device: str = 'cuda',
+        use_skip=False,
     ) -> None:
         super().__init__()
 
+        self.use_skip = use_skip
+
         if not block_setting:
             raise ValueError("The block_setting should not be empty")
         elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
@@ -259,10 +262,13 @@ class ConvNeXtReverseBase2(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
-        x, features = x
-        x = self.first(x)
+        input, features = x
+        x = self.first(input)
         for i in range(5 - 2):
-            x = self.stages[i](x + features[2-i])
+            if self.use_skip:
+                x = self.stages[i](torch.concat((x, features[2-i]), dim=1))
+            else:
+                x = self.stages[i](x)
             x = self.upsamples[i](x)
         
         x = self.stages[-2](x)