Skip to content

Instantly share code, notes, and snippets.

@agastidukare
Created April 1, 2020 02:14
Show Gist options
  • Save agastidukare/09b63738900870fa97df5c6c6a96942b to your computer and use it in GitHub Desktop.
Save agastidukare/09b63738900870fa97df5c6c6a96942b to your computer and use it in GitHub Desktop.
init_random_params, predict = stax.serial(
stax.Conv(64, (7,7), padding='SAME'),
stax.Relu,
stax.Conv(32, (4, 4), padding='SAME'),
stax.Relu,
stax.MaxPool((3, 3)),
stax.Flatten,
stax.Dense(128),
stax.Relu,
stax.Dense(10),
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment