Update linear.py

ASEM000 · Aug 1, 2023 · 6115e18 · 6115e18
1 parent 3984627
commit 6115e18
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/serket/nn/linear.py b/serket/nn/linear.py
@@ -566,6 +566,7 @@ class MLP(sk.TreeClass):
         >>> lazy_mlp = sk.nn.MLP(None, 1, num_hidden_layers=2, hidden_size=10)
         >>> _, materialized_mlp = lazy_mlp.at['__call__'](jnp.ones([1, 10]))
         >>> materialized_mlp.layers[0].in_features
+        (10,)
     """
 
     def __init__(