add cvae training

models/convnext_orig.py

@@ -66,6 +66,7 @@ class ConvNeXtBase(nn.Module):
         kernel_size: int = 7,
         is_label: bool = False,
         nr_of_labels: Optional[int] = None,
+        input_features: int = 3,
     ) -> None:
         super().__init__()
 
@@ -85,7 +86,7 @@ class ConvNeXtBase(nn.Module):
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
         self.stem = Conv2dNormActivation(
-                3 if not is_label else nr_of_labels,
+                input_features if not is_label else nr_of_labels,
                 firstconv_output_channels,
                 kernel_size=4,
                 stride=4,
@@ -229,6 +230,7 @@ def _convnext(
     is_label: bool,
     nr_of_labels: Optional[int],
     kernel_size: int,
+    input_features: int,
     device='cuda',
     **kwargs: Any,
 ) -> ConvNeXtBase:
@@ -236,17 +238,17 @@ def _convnext(
     if reversed:
         model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
     else:
-        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device, **kwargs)
+        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, input_features=input_features, device=device, **kwargs)
 
 
     return model
 
 
 class ConvNeXt(nn.Module):
-    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, is_label=False, nr_of_labels=None, kernel_size=7, device='cuda', **kwargs) -> None:
+    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, is_label=False, nr_of_labels=None, kernel_size=7, input_features=3, device='cuda', **kwargs) -> None:
         super().__init__()
 
-        self.model = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, device=device, **kwargs)
+        self.model: nn.Module = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, input_features=input_features, device=device, **kwargs)
 
     def forward(self, x: Tensor) -> Tensor:
         x = self.model(x)