add-bn

src/eva/core/models/networks/mlp.py

@@ -17,6 +17,7 @@ def __init__(
         hidden_activation_fn: Type[torch.nn.Module] | None = nn.ReLU,
         output_activation_fn: Type[torch.nn.Module] | None = None,
         dropout: float = 0.0,
+        use_batch_norm: bool = False,
     ) -> None:
         """Initializes the MLP.
 
@@ -36,6 +37,7 @@ def __init__(
         self.hidden_activation_fn = hidden_activation_fn
         self.output_activation_fn = output_activation_fn
         self.dropout = dropout
+        self.use_batch_norm = use_batch_norm
 
         self._network = self._build_network()
 
@@ -45,6 +47,8 @@ def _build_network(self) -> nn.Sequential:
         prev_size = self.input_size
         for size in self.hidden_layer_sizes:
             layers.append(nn.Linear(prev_size, size))
+            if self.use_batch_norm:
+                layers.append(nn.BatchNorm1d(size))
             if self.hidden_activation_fn is not None:
                 layers.append(self.hidden_activation_fn())
             if self.dropout > 0: