curvature training

models/convnext_orig.py

@@ -64,6 +64,8 @@ class ConvNeXtBase(nn.Module):
         norm_layer: Optional[Callable[..., nn.Module]] = None,
         device: str = 'cuda',
         kernel_size: int = 7,
+        is_label: bool = False,
+        nr_of_labels: Optional[int] = None,
     ) -> None:
         super().__init__()
 
@@ -83,7 +85,7 @@ class ConvNeXtBase(nn.Module):
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
         self.stem = Conv2dNormActivation(
-                3,
+                3 if not is_label else nr_of_labels,
                 firstconv_output_channels,
                 kernel_size=4,
                 stride=4,
@@ -234,7 +236,7 @@ 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, 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, device=device, **kwargs)
 
 
     return model