Change default checkpoint saving dtype to float32 instead of bfloat16…

…. Saving PiperOrigin-RevId: 364201025
tensorflow · Mar 21, 2021 · 4e07d5e · 4e07d5e
1 parent 5d16bca
commit 4e07d5e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mesh_tensorflow/transformer/utils.py b/mesh_tensorflow/transformer/utils.py
@@ -87,7 +87,7 @@ def parse_gin_defaults_and_flags(skip_unknown=False, finalize_config=True):
 #  this stupid VariableDtype class and stop passing it all over creation.
 @gin.configurable
 def get_variable_dtype(
-    master_dtype=tf.float32,
+    master_dtype=tf.bfloat16,
     slice_dtype=tf.float32,
     activation_dtype=tf.float32):
   """Datatypes to use for the run.