inspect vae structure

prishajain1 · prishajain1 · commit b474e1ab5900 · 2026-02-23T10:51:17.000+05:30
diff --git a/inspect_vae_structure.py b/inspect_vae_structure.py
@@ -18,7 +18,8 @@ def inspect_structure():
         decoder_inject_noise=(False, False),
         downsample_type=("spatial", "temporal"),
         upsample_residual=(True, True),
-        upsample_factor=(2, 2)
+        upsample_factor=(2, 2),
+        rngs=nnx.Rngs(0)
     )
     
     state = nnx.state(model)

Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,8 @@ def inspect_structure():`
`18`	`18`	`decoder_inject_noise=(False, False),`
`19`	`19`	`downsample_type=("spatial", "temporal"),`
`20`	`20`	`upsample_residual=(True, True),`
`21`		`- upsample_factor=(2, 2)`
	`21`	`+ upsample_factor=(2, 2),`
	`22`	`+ rngs=nnx.Rngs(0)`
`22`	`23`	`)`
`23`	`24`
`24`	`25`	`state = nnx.state(model)`