Modify `enable_pinned_host_transfer` option to be provided once for the entire pytree, since it's not really reasonable to customize this on a per-array level.

google / orbax

Orbax provides common checkpointing and persistence utilities for JAX users

https://orbax.readthedocs.io/

Apache License 2.0

305 stars 36 forks source link

Modify `enable_pinned_host_transfer` option to be provided once for the entire pytree, since it's not really reasonable to customize this on a per-array level. #1239

Closed copybara-service[bot] closed 1 month ago

copybara-service[bot] commented 1 month ago

Modify enable_pinned_host_transfer option to be provided once for the entire pytree, since it's not really reasonable to customize this on a per-array level.