cahya's picture
reorder "import jax" to avoid the stuck on importing it, change the dataset
b25d188