Skip to content

Instantly share code, notes, and snippets.

@agastidukare
Created April 1, 2020 02:18
Show Gist options
  • Save agastidukare/eab0e10cc32fe9713a8e104c058e1867 to your computer and use it in GitHub Desktop.
Save agastidukare/eab0e10cc32fe9713a8e104c058e1867 to your computer and use it in GitHub Desktop.
learning_rate = 0.14
opt_init, opt_update, get_params = optimizers.sgd(learning_rate)
@jit
def update(_, i, opt_state, batch):
params = get_params(opt_state)
return opt_update(i, grad(loss)(params, batch), opt_state)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment