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)

与 save_only_these_names 相同，但将中间结果卸载到 CPU 内存而不是重新计算。