Skip to content

Instantly share code, notes, and snippets.

@agastidukare
Created April 1, 2020 02:25
Show Gist options
  • Save agastidukare/8eaf256432a5ca1699cea13b2ef8f14b to your computer and use it in GitHub Desktop.
Save agastidukare/8eaf256432a5ca1699cea13b2ef8f14b to your computer and use it in GitHub Desktop.
num_epochs = 1
key = random.PRNGKey(123)
_, init_params = init_random_params(key, (-1, 28, 28, 1))
opt_state = opt_init(init_params)
itercount = itertools.count()
for _ in range(num_batches):
opt_state= update(key, next(itercount), opt_state, shape_as_image(*next(batches)))
params = get_params(opt_state)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment