def METHOD_NAME(rng):
  quick_draw = QuickDraw()
  params = quick_draw.init(rng, jnp.ones((1, 28, 28, 1)))['params']
  tx = optax.adam(learning_rate=0.001, b1=0.9, b2=0.999)
  return train_state.TrainState.create(
    apply_fn=quick_draw.apply, params=params, tx=tx)