# Buggy code
def decode(
    self,
    inputs,
    sequence_length,
    vocab_size=None,
    initial_state=None,
    sampling_probability=None,
    embedding=None,
    output_layer=None,
    mode=tf.estimator.ModeKeys.TRAIN,
    memory=None,
    memory_sequence_length=None,
):
    _ = memory
    _ = memory_sequence_length

    batch_size = tf.shape(inputs)[0]

    if sampling_probability is not None and (
        tf.contrib.framework.is_tensor(sampling_probability)
        or sampling_probability > 0.0
    ):
        if embedding is None:
            raise ValueError(
                "embedding argument must be set when using scheduled sampling"
            )

        tf.summary.scalar("sampling_probability", sampling_probability)
        helper = tf.contrib.seq2seq.ScheduledEmbeddingTrainingHelper(
            inputs, sequence_length, embedding, sampling_probability
        )
    else:
        helper = tf.contrib.seq2seq.TrainingHelper(inputs, sequence_length)

    cell, initial_state = self._build_cell(
        mode,
        batch_size,
        initial_state=initial_state,
        memory=memory,
        memory_sequence_length=memory_sequence_length,
        dtype=inputs.dtype,
    )

    if output_layer is None:
        output_layer = build_output_layer(
            self.num_units, vocab_size, dtype=inputs.dtype
        )

    # With TrainingHelper, project all timesteps at once.
    fused_projection = isinstance(helper, tf.contrib.seq2seq.TrainingHelper)

    decoder = tf.contrib.seq2seq.BasicDecoder(
        cell,
        helper,
        initial_state,
        output_layer=output_layer if not fused_projection else None,
    )

    outputs, state, length = tf.contrib.seq2seq.dynamic_decode(decoder)

    if fused_projection and output_layer is not None:
        logits = output_layer(outputs.rnn_output)
    else:
        logits = outputs.rnn_output
    # Make sure outputs have the same time_dim as inputs
    inputs_len = tf.shape(inputs)[1]
    logits = align_in_time(logits, inputs_len)

    return (logits, state, length)


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
Traceback (most recent call last):
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1323, in _do_call
return fn(*args)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1302, in _run_fn
status, run_metadata)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/errors_impl.py", line 473, in __exit__
c_api.TF_GetCode(self.status.status))
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0] = 319 is not in [0, 167)
[[Node: seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/decoder_1/embedding_lookup = Gather[Tindices=DT_INT3
2, Tparams=DT_FLOAT, _class=["loc:@seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/decoder_1/convert_gradient_to_tens
or_cc661786"], validate_indices=true, _device="/job:localhost/replica:0/task:0/device:CPU:0"](seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrain
ingHelperNextInputs/cond/decoder_1/convert_gradient_to_tensor_cc661786, seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/co
nd/GatherNd)]]
[[Node: seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/Greater/Switch/_525 = _HostRecv[client_terminated=fa
lse, recv_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device_incarnation=1, tensor_name="edge_7992_seq2seq/pa
rallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/Greater/Switch", tensor_type=DT_INT32, _device="/job:localhost/replica:0/task:0
/device:GPU:0"](^_cloopseq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/TrainingHelperNextInputs/GreaterEqual/_132)]]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/hltsrv0/rdessi/sw/anaconda3/bin/onmt-main", line 11, in <module>
load_entry_point('OpenNMT-tf==1.5.0', 'console_scripts', 'onmt-main')()
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/bin/main.py", line 133, in main
runner.train_and_evaluate()
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/runner.py", line 148, in train_and_evaluate
tf.estimator.train_and_evaluate(self._estimator, train_spec, eval_spec)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/training.py", line 432, in train_and_evaluate
executor.run_local()
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/training.py", line 611, in run_local
hooks=train_hooks)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 302, in train
loss = self._train_model(input_fn, hooks, saving_listeners)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 783, in _train_model
_, loss = mon_sess.run([estimator_spec.train_op, estimator_spec.loss])
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 521, in run
run_metadata=run_metadata)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 892, in run
run_metadata=run_metadata)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 967, in run
raise six.reraise(*original_exc_info)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/six.py", line 693, in reraise
raise value
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 952, in run
return self._sess.run(*args, **kwargs)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1024, in run
run_metadata=run_metadata)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 827, in run
return self._sess.run(*args, **kwargs)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 889, in run
run_metadata_ptr)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1120, in _run
feed_dict_tensor, options, run_metadata)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1317, in _do_run
options, run_metadata)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1336, in _do_call
raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0] = 319 is not in [0, 167)
[[Node: seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/decoder_1/embedding_lookup = Gather[Tindices=DT_INT32, Tparams=DT_FLOAT, _class=["loc:@seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/decoder_1/convert_gradient_to_tensor_cc661786"], validate_indices=true, _device="/job:localhost/replica:0/task:0/device:CPU:0"](seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/decoder_1/convert_gradient_to_tensor_cc661786, seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/GatherNd)]]
[[Node: seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/Greater/Switch/_525 = _HostRecv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device_incarnation=1, tensor_name="edge_7992_seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/Greater/Switch", tensor_type=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](^_cloopseq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/TrainingHelperNextInputs/GreaterEqual/_132)]]

Caused by op 'seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/decoder_1/embedding_lookup', defined at:
File "/hltsrv0/rdessi/sw/anaconda3/bin/onmt-main", line 11, in <module>
load_entry_point('OpenNMT-tf==1.5.0', 'console_scripts', 'onmt-main')()
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/bin/main.py", line 133, in main
runner.train_and_evaluate()
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/runner.py", line 148, in train_and_evaluate
tf.estimator.train_and_evaluate(self._estimator, train_spec, eval_spec)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/training.py", line 432, in train_and_evaluate
executor.run_local()
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/training.py", line 611, in run_local
hooks=train_hooks)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 302, in train
loss = self._train_model(input_fn, hooks, saving_listeners)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 711, in _train_model
features, labels, model_fn_lib.ModeKeys.TRAIN, self.config)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 694, in _call_model_fn
model_fn_results = self._model_fn(features=features, **kwargs)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/models/model.py", line 103, in _model_fn
_loss_op, features_shards, labels_shards, params, mode, config)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/utils/parallel.py", line 148, in __call__
outputs.append(funs[i](*args[i], **kwargs[i]))
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/models/model.py", line 66, in _loss_op
logits, _ = self._build(features, labels, params, mode, config=config)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/models/sequence_to_sequence.py", line 144, in _build
memory_sequence_length=encoder_sequence_length)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/decoders/rnn_decoder.py", line 130, in decode
outputs, state, length = tf.contrib.seq2seq.dynamic_decode(decoder)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/seq2seq/python/ops/decoder.py", line 286, in dynamic_decode
swap_memory=swap_memory)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2816, in while_loop
result = loop_context.BuildLoop(cond, body, loop_vars, shape_invariants)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2640, in BuildLoop
pred, body, original_loop_vars, loop_vars, shape_invariants)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2590, in _BuildLoop
body_result = body(*packed_vars_for_body)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/seq2seq/python/ops/decoder.py", line 234, in body
decoder_finished) = decoder.step(time, inputs, state)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py", line 147, in step
sample_ids=sample_ids)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/seq2seq/python/ops/helper.py", line 342, in next_inputs
all_finished, lambda: base_next_inputs, maybe_sample)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/util/deprecation.py", line 316, in new_func
return func(*args, **kwargs)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 1864, in cond
orig_res_f, res_f = context_f.BuildCondBranch(false_fn)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 1725, in BuildCondBranch
original_result = fn()
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/seq2seq/python/ops/helper.py", line 331, in maybe_sample
sampled_next_inputs = self._embedding_fn(sample_ids_sampling)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/models/sequence_to_sequence.py", line 103, in _target_embedding_fn
return self.target_inputter.transform(ids, mode=mode)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/inputters/text_inputter.py", line 392, in transform
outputs = embedding_lookup(embeddings, inputs)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/OpenNMT_tf-1.5.0-py3.6.egg/opennmt/layers/common.py", line 30, in embedding_lookup
return tf.nn.embedding_lookup(params, ids)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/embedding_ops.py", line 328, in embedding_lookup
transform_fn=None)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/embedding_ops.py", line 150, in _embedding_lookup_and_transform
result = _clip(_gather(params[0], ids, name=name), ids, max_norm)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/embedding_ops.py", line 54, in _gather
return array_ops.gather(params, ids, name=name)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/array_ops.py", line 2486, in gather
params, indices, validate_indices=validate_indices, name=name)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/gen_array_ops.py", line 1834, in gather
validate_indices=validate_indices, name=name)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py", line 787, in _apply_op_helper
op_def=op_def)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 2956, in create_op
op_def=op_def)
File "/hltsrv0/rdessi/sw/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 1470, in __init__
self._traceback = self._graph._extract_stack()  # pylint: disable=protected-access

InvalidArgumentError (see above for traceback): indices[0] = 319 is not in [0, 167)
[[Node: seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/decoder_1/embedding_lookup = Gather[Tindices=DT_INT32, Tparams=DT_FLOAT, _class=["loc:@seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/decoder_1/convert_gradient_to_tensor_cc661786"], validate_indices=true, _device="/job:localhost/replica:0/task:0/device:CPU:0"](seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/decoder_1/convert_gradient_to_tensor_cc661786, seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/GatherNd)]]
[[Node: seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/Greater/Switch/_525 = _HostRecv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device_incarnation=1, tensor_name="edge_7992_seq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/cond/Greater/Switch", tensor_type=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](^_cloopseq2seq/parallel_0/seq2seq/decoder/decoder_1/while/BasicDecoderStep/ScheduledEmbeddingTrainingHelperNextInputs/TrainingHelperNextInputs/GreaterEqual/_132)]]

# Fixed code
def decode(
    self,
    inputs,
    sequence_length,
    vocab_size=None,
    initial_state=None,
    sampling_probability=None,
    embedding=None,
    output_layer=None,
    mode=tf.estimator.ModeKeys.TRAIN,
    memory=None,
    memory_sequence_length=None,
):
    _ = memory
    _ = memory_sequence_length

    batch_size = tf.shape(inputs)[0]

    if sampling_probability is not None and (
        tf.contrib.framework.is_tensor(sampling_probability)
        or sampling_probability > 0.0
    ):
        if embedding is None:
            raise ValueError(
                "embedding argument must be set when using scheduled sampling"
            )

        tf.summary.scalar("sampling_probability", sampling_probability)
        helper = tf.contrib.seq2seq.ScheduledEmbeddingTrainingHelper(
            inputs, sequence_length, embedding, sampling_probability
        )
        fused_projection = False
    else:
        helper = tf.contrib.seq2seq.TrainingHelper(inputs, sequence_length)
        fused_projection = True  # With TrainingHelper, project all timesteps at once.

    cell, initial_state = self._build_cell(
        mode,
        batch_size,
        initial_state=initial_state,
        memory=memory,
        memory_sequence_length=memory_sequence_length,
        dtype=inputs.dtype,
    )

    if output_layer is None:
        output_layer = build_output_layer(
            self.num_units, vocab_size, dtype=inputs.dtype
        )

    decoder = tf.contrib.seq2seq.BasicDecoder(
        cell,
        helper,
        initial_state,
        output_layer=output_layer if not fused_projection else None,
    )

    outputs, state, length = tf.contrib.seq2seq.dynamic_decode(decoder)

    if fused_projection and output_layer is not None:
        logits = output_layer(outputs.rnn_output)
    else:
        logits = outputs.rnn_output
    # Make sure outputs have the same time_dim as inputs
    inputs_len = tf.shape(inputs)[1]
    logits = align_in_time(logits, inputs_len)

    return (logits, state, length)