File "/lib/python3.5/dist-packages/tensorflow_estimator/python/estimator/estimator.py", line 1149, in _call_model_fn
    model_fn_results = self._model_fn(features=features, **kwargs)
  File "/lib/python3.5/dist-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3159, in _model_fn
    _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn))
  File "/lib/python3.5/dist-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3604, in _train_on_tpu_system
    device_assignment=ctx.device_assignment)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/tpu/tpu.py", line 1277, in split_compile_and_shard
    name=name)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/tpu/tpu.py", line 992, in split_compile_and_replicate
    outputs = computation(*computation_inputs)
  File "/lib/python3.5/dist-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3589, in multi_tpu_train_steps_on_single_shard
    inputs=[0, _INITIAL_LOSS])
  File "/lib/python3.5/dist-packages/tensorflow_core/python/tpu/training_loop.py", line 178, in while_loop
    condition_wrapper, body_wrapper, inputs, name="", parallel_iterations=1)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/ops/control_flow_ops.py", line 2753, in while_loop
    return_same_structure)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/ops/control_flow_ops.py", line 2245, in BuildLoop
    pred, body, original_loop_vars, loop_vars, shape_invariants)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/ops/control_flow_ops.py", line 2170, in _BuildLoop
    body_result = body(*packed_vars_for_body)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/tpu/training_loop.py", line 121, in body_wrapper
    outputs = body(*(inputs + dequeue_ops))
  File "/lib/python3.5/dist-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3588, in <lambda>
    lambda i, loss: [i + 1, single_tpu_train_step(i)],
  File "/lib/python3.5/dist-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 1715, in train_step
    self._call_model_fn(features, labels))
  File "/lib/python3.5/dist-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 1994, in _call_model_fn
    estimator_spec = self._model_fn(features=features, **kwargs)
  File "/lib/python3.5/dist-packages/mesh_tensorflow/transformer/utils.py", line 530, in my_model_fn
    lowering = mtf.Lowering(graph, {mesh: mesh_impl}, autostack=autostack)
  File "/lib/python3.5/dist-packages/mesh_tensorflow/ops.py", line 653, in __init__
    op.lower(self)
  File "/lib/python3.5/dist-packages/mesh_tensorflow/ops.py", line 3842, in lower
    sv = mesh_impl.LaidOutVariable(self, mesh_impl)
  File "/lib/python3.5/dist-packages/mesh_tensorflow/simd_mesh_impl.py", line 202, in __init__
    variable.get_master(), shape, slices, slice_shape)
  File "/lib/python3.5/dist-packages/mesh_tensorflow/simd_mesh_impl.py", line 228, in _gen_copy_master_to_slices_op
    master_variable = tf.cast(master_variable, slices[0].dtype)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/util/dispatch.py", line 180, in wrapper
    return target(*args, **kwargs)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/ops/math_ops.py", line 704, in cast
    x = gen_math_ops.cast(x, base_type, name=name)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/ops/gen_math_ops.py", line 2219, in cast
    "Cast", x=x, DstT=DstT, Truncate=Truncate, name=name)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/framework/op_def_library.py", line 794, in _apply_op_helper
    op_def=op_def)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/framework/ops.py", line 3357, in create_op
    attrs, op_def, compute_device)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/framework/ops.py", line 3426, in _create_op_internal
    op_def=op_def)
  File "/lib/python3.5/dist-packages/tensorflow_core/python/framework/ops.py", line 1748, in __init__
    self._traceback = tf_stack.extract_stack()

  In call to configurable 'run' (<function run at 0x7f0c1e948bf8>)