Change default checkpoint saving dtype to float32 instead of bfloat16…

…. Saving PiperOrigin-RevId: 363758862
tensorflow · Mar 18, 2021 · 5d16bca · 5d16bca
1 parent 441ff47
commit 5d16bca
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.bfloat16,
+    master_dtype=tf.float32,
     slice_dtype=tf.float32,
     activation_dtype=tf.float32):
   """Datatypes to use for the run.