Always compute the softmax in full precision
Related to
@@ -14,6 +14,9 @@ def _smooth_one_hot_labels(logits, labels, label_smoothing):
dtype=logits.dtype)
def _softmax_cross_entropy(logits, labels, label_smoothing, mode):
+ # Computes the softmax in full precision.
+ if logits.dtype.base_dtype != tf.float32:
+ logits = tf.cast(logits, tf.float32)
if mode == tf.estimator.ModeKeys.TRAIN and label_smoothing > 0.0:
smoothed_labels = _smooth_one_hot_labels(logits, labels, label_smoothing)
if hasattr(tf.nn, "softmax_cross_entropy_with_logits_v2"):
