Closed this issue 4 years ago · 1 comments
Colab now runs jax v0.2, and apparently some of our code is no longer working. This PR is to track bugs and fixes to update to recent jax version.
ok this got fixed by #67 closing this issue