diff --git a/src/MLPs/centralized.py b/src/MLPs/centralized.py index 70abbec..14e3a22 100644 --- a/src/MLPs/centralized.py +++ b/src/MLPs/centralized.py @@ -6,6 +6,7 @@ from flax.linen.initializers import constant, orthogonal class Network(nn.Module): hidden_size: int = 256 + @nn.compact def __call__(self, x): # x shape: (batch, obs_dim)