modify SGD training to avoid using generators

Open murphyk opened this issue 3 years ago • 1 comments

Currently, run_sgd uses generators to create a stream of minibatches, but this causes issues with jit. It may be better to refactor the code to work in terms of epochs, and then shuffle at the start, similar to this flax mnist example

Nov 09 '22 18:11 murphyk

See also https://github.com/probml/dynamax/issues/130

Nov 14 '22 06:11 murphyk