use jnp.stack instead of np.stack - SemanticDiff

Commit

157 days ago

use jnp.stack instead of np.stack

References

#4991 - Use `jnp.stack` instead of `np.stack` in `flax.training.common_utils.stack_forest`

Author

vfdev-5

vfdev-5

Parents

Loading