View source on GitHub |
Saves name based Tensor into a Checkpoint.
tf.experimental.dtensor.name_based_save(
mesh: tf.experimental.dtensor.Mesh
,
checkpoint_prefix: Union[str, tf.Tensor
],
name_tensor_dict: Dict[str, Union[ops.Tensor, tf_variables.Variable]]
)
The function prepares the input dictionary to the format of a sharded_save
,
so that it can take advantage of DTensor SPMD based distributed save.
Same as restore, the function only supports saving on the single mesh.