Document context_model using SumContextEncoder
Browse files
README.md
CHANGED
|
@@ -51,3 +51,16 @@ checkpointer = Checkpointer(
|
|
| 51 |
)
|
| 52 |
gen_model = nnx.merge(graphdef, checkpointer.restore_ema(params))
|
| 53 |
```
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 51 |
)
|
| 52 |
gen_model = nnx.merge(graphdef, checkpointer.restore_ema(params))
|
| 53 |
```
|
| 54 |
+
|
| 55 |
+
`restore_ema` reads only `ema_params` and `extra`, and applies the same bias
|
| 56 |
+
correction the training loop uses for evaluation.
|
| 57 |
+
|
| 58 |
+
## Reproducing
|
| 59 |
+
|
| 60 |
+
```bash
|
| 61 |
+
uv run python tutorials/scripts/train_conditioning_checkpoints.py
|
| 62 |
+
```
|
| 63 |
+
|
| 64 |
+
The script mirrors the notebook's model definitions and PRNG chain, so it
|
| 65 |
+
reproduces these exact weights. A checkpoint pins the parameter structure: if a
|
| 66 |
+
notebook's network changes, re-run the script and re-upload.
|