fix curvature model

convnext/models/convnext_tf.py

@@ -348,7 +348,7 @@ class ConvNeXtReverseLabels(tf.keras.Model):
             LayerNormalization(epsilon=1e-6)]
         )
         self.downsample_layers.append(stem)
-        for i in range(3):
+        for i in range(len(depths) - 1):
             downsample_layer = tf.keras.Sequential([
                 LayerNormalization(epsilon=1e-6),
                 Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
@@ -358,7 +358,7 @@ class ConvNeXtReverseLabels(tf.keras.Model):
         self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
         dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
         cur = 0
-        for i in range(4):
+        for i in range(len(depths)):
             stage = tf.keras.Sequential(
                 [Block(dim=dims[i], drop_path=dp_rates[cur + j],
                        layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
@@ -369,7 +369,7 @@ class ConvNeXtReverseLabels(tf.keras.Model):
         self.last = Dense(nr_of_classes, activation='softmax')
 
     def call(self, x):
-        for i in reversed(range(4)):
+        for i in reversed(len(depths)):
             x = self.downsample_layers[i](x)
             x = self.stages[i](x)
         x = self.last(x)