attribute

yallup · Jun 24, 2024 · 0d42f30 · 0d42f30
1 parent 3be37bf
commit 0d42f30
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/clax/clax.py b/clax/clax.py
@@ -37,6 +37,7 @@ def __init__(self, n=1, **kwargs):
             seed: int: Random seed.
         """
         self.rng = random.PRNGKey(kwargs.get("seed", 2024))
+        self.n = n
         self.network = Network(n_out=n)
         self.state = None