flax
Rewritten flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate using jax.device_put
#5101

Open

Commits

Rewritten flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate using jax.device_put

vfdev-5 committed 84 days ago

Loading