Merge pull request #3703 from levskaya:rematdocfix

PiperOrigin-RevId: 608683576
google · Feb 20, 2024 · eca5dd6 · eca5dd6
2 parents efbe705 + 5c7fb4c
commit eca5dd6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
@@ -971,12 +971,12 @@ def checkpoint(
     >>> import flax.linen as nn
     ...
     >>> class CheckpointedMLP(nn.Module):
+    ...   @nn.checkpoint
     ...   @nn.compact
     ...   def __call__(self, x):
-    ...     CheckpointDense = nn.checkpoint(nn.Dense)
-    ...     x = CheckpointDense(128)(x)
+    ...     x = nn.Dense(128)(x)
     ...     x = nn.relu(x)
-    ...     x = CheckpointDense(1)(x)
+    ...     x = nn.Dense(1)(x)
     ...     return x
     ...
     >>> model = CheckpointedMLP()