1
Fork 0

fix: uniform generic network naming

This commit is contained in:
cedric 2026-04-02 18:18:10 +00:00
parent 40c7e32344
commit 79221372df
2 changed files with 10 additions and 4 deletions

View file

@ -9,7 +9,7 @@ from flax.linen.initializers import constant, orthogonal
# semi generic so we can easily make a config for it in experiments
class SemiGenericNetwork(nn.Module):
class GenericDenseLayersWithActivation(nn.Module):
layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
activation: Callable = nn.tanh