jax.checkpoint_policies.offload_dot_with_no_batch_dims

checkpoint_policies.offload_dot_with_no_batch_dims(offload_dst)

与 dots_with_no_batch_dims_saveable 相同，但会将数据卸载到 CPU 内存而不是重新计算。

这是Transformer的一个有用启发式方法。