tf.experimental.dtensor.enable_save_as_bf16
Allows float32 DVariables to be checkpointed and restored as bfloat16.
tf.experimental.dtensor.enable_save_as_bf16(
variables: List[tf.Variable]
)
The method only affects the DVariable part inside the model and leaves non-DTensor Variables/Tensors untouched.
| Args | |
|---|---|
variables |
A list of tf.Variable to be enabled with bfloat16 save/restore. Only has effect on DTensor Variables as they go through d_variables with DTensor Specific logis. |
© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/experimental/dtensor/enable_save_as_bf16
