message: Update calling of expand_dims with axis
  This fix updates calling of `expand_dims` with `dim -> axis`
  as the `dim=` in `tf.expand_dims` has been deprecated and was
  generating unnecessary warnings.
diff: @@ -563,10 +563,10 @@ def _mean_squared_loss(labels, logits, weights=None):
  labels = ops.convert_to_tensor(labels)
  # To prevent broadcasting inside "-".
  if len(labels.get_shape()) == 1:
  - labels = array_ops.expand_dims(labels, dim=(1,))
  + labels = array_ops.expand_dims(labels, axis=(1,))
  # TODO(zakaria): make sure it does not recreate the broadcast bug.
  if len(logits.get_shape()) == 1:
  - logits = array_ops.expand_dims(logits, dim=(1,))
  + logits = array_ops.expand_dims(logits, axis=(1,))
  logits.get_shape().assert_is_compatible_with(labels.get_shape())
  loss = math_ops.square(logits - math_ops.to_float(labels), name=name)
  return _compute_weighted_loss(loss, weights)
  