Add a classifier option to convnets

ziatdinovmax · Jan 9, 2025 · 0b3de08 · 0b3de08
1 parent 4d7db0e
commit 0b3de08
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/neurobayes/flax_nets/convnet.py b/neurobayes/flax_nets/convnet.py
@@ -28,6 +28,7 @@ class FlaxConvNet(nn.Module):
     target_dim: int
     activation: str = 'tanh'
     kernel_size: Union[int, Tuple[int, ...]] = 3
+    classification: bool = False  # Explicit flag for classification tasks
 
     @nn.compact
     def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
@@ -51,7 +52,8 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
         x = FlaxMLP(
             hidden_dims=self.fc_layers,
             target_dim=self.target_dim,
-            activation=self.activation)(x)
+            activation=self.activation,
+            classification=self.classification)(x)
 
         return x