jax.checkpoint_policies.save_and_offload_only_these_names

Contents

jax.checkpoint_policies.save_and_offload_only_these_names#

checkpoint_policies.save_and_offload_only_these_names(*, names_which_can_be_offloaded, offload_src, offload_dst)[source]#: Same as save_only_these_names, but offload to CPU memory instead of recomputing.