search

cgarciae / treex

A Pytree Module system for Deep Learning in JAX
https://cgarciae.github.io/treex/
MIT License
215 stars 17 forks source link

Fix assert key jnp.ndarray error #10

Closed lkhphuc closed 2 years ago

lkhphuc commented 2 years ago

I'm not very familiar with the code base but this at least make the example run. Hope it help. Should fix #9