message: Getting initialized value from the same device as the original variable,
  when creating EMA variables.
diff: @@ -453,9 +453,11 @@ class ExponentialMovingAverage(object):
  # tensors, we rely on the existing device allocation mechanism.
  with ops.init_scope():
  if isinstance(var, variables.Variable):
  + with ops.device(var.device):
  + initialized_value = var.initialized_value()
  avg = slot_creator.create_slot(
  var,
  - var.initialized_value(),
  + initialized_value,
  self.name,
  colocate_with_primary=True,
  copy_xla_sharding=True)
  