Closed qlzh727 closed 1 year ago
Patch coverage: 100.00%
and project coverage change: +5.07%
:tada:
Comparison is base (
67722d7
) 70.92% compared to head (d32e52e
) 76.00%. Report is 2 commits behind head on main.
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
The unit test failed after we change to use tree.map_structure(), probably it handles None differently than jax.tree_map(). Should I still prefer tree.map_structure()?
The unit test failed after we change to use tree.map_structure(), probably it handles None differently than jax.tree_map(). Should I still prefer tree.map_structure()?
Add an additional None handle for this.
The current eager_build will use eager tensor to run the forward path of the model to trigger variable creation, but this will bring non-trivial memory consumption from the forward path, and will cause OOM when the model are large, and designed for distribution computation.
This PR will use the kerasTensor instead, which doesn't consume any memory.
Verified with local colab for memory consumption.