# 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

    return (logits, state, length)


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
Traceback (most recent call last):
File "/root/tf-1.4/bin/onmt-main", line 11, in <module>
sys.exit(main())
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/bin/main.py", line 120, in main
runner.train_and_evaluate()
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/runner.py", line 141, in train_and_evaluate
tf.estimator.train_and_evaluate(self._estimator, train_spec, eval_spec)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/estimator/training.py", line 430, in train_and_evaluate
executor.run_local()
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/estimator/training.py", line 609, in run_local
hooks=train_hooks)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/estimator/estimator.py", line 302, in train
loss = self._train_model(input_fn, hooks, saving_listeners)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/estimator/estimator.py", line 783, in _train_model
_, loss = mon_sess.run([estimator_spec.train_op, estimator_spec.loss])
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/training/monitored_session.py", line 521, in run
run_metadata=run_metadata)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/training/monitored_session.py", line 892, in run
run_metadata=run_metadata)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/training/monitored_session.py", line 967, in run
raise six.reraise(*original_exc_info)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/training/monitored_session.py", line 952, in run
return self._sess.run(*args, **kwargs)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/training/monitored_session.py", line 1024, in run
run_metadata=run_metadata)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/training/monitored_session.py", line 827, in run
return self._sess.run(*args, **kwargs)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 889, in run
run_metadata_ptr)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 1120, in _run
feed_dict_tensor, options, run_metadata)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 1317, in _do_run
options, run_metadata)
File "/root/tf-1.4/local/lib/python2.7/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: logits and labels must be same size: logits_size=[196,101] labels_size=[203,101]
[[Node: seq2seq/parallel_1/seq2seq/SoftmaxCrossEntropyWithLogits = SoftmaxCrossEntropyWithLogits[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"](seq2seq/parallel_1/seq2seq/Reshape, seq2seq/parallel_1/seq2seq/Reshape_1)]]

Caused by op u'seq2seq/parallel_1/seq2seq/SoftmaxCrossEntropyWithLogits', defined at:
File "/root/tf-1.4/bin/onmt-main", line 11, in <module>
sys.exit(main())
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/bin/main.py", line 120, in main
runner.train_and_evaluate()
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/runner.py", line 141, in train_and_evaluate
tf.estimator.train_and_evaluate(self._estimator, train_spec, eval_spec)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/estimator/training.py", line 430, in train_and_evaluate
executor.run_local()
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/estimator/training.py", line 609, in run_local
hooks=train_hooks)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/estimator/estimator.py", line 302, in train
loss = self._train_model(input_fn, hooks, saving_listeners)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/estimator/estimator.py", line 711, in _train_model
features, labels, model_fn_lib.ModeKeys.TRAIN, self.config)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/estimator/estimator.py", line 694, in _call_model_fn
model_fn_results = self._model_fn(features=features, **kwargs)
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/models/model.py", line 88, in _model_fn
_loss_op, features_shards, labels_shards, params, mode, config)
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/utils/parallel.py", line 150, in __call__
outputs.append(funs[i](*args[i], **kwargs[i]))
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/models/model.py", line 52, in _loss_op
return self._compute_loss(features, labels, logits, params, mode)
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/models/sequence_to_sequence.py", line 229, in _compute_loss
mode=mode)
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/utils/losses.py", line 51, in cross_entropy_sequence_loss
cross_entropy = _softmax_cross_entropy(logits, labels, label_smoothing, mode)
File "/root/tf-1.4/local/lib/python2.7/site-packages/opennmt/utils/losses.py", line 24, in _softmax_cross_entropy
logits=logits, labels=smoothed_labels)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/ops/nn_ops.py", line 1783, in softmax_cross_entropy_with_logits
precise_logits, labels, name=name)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/ops/gen_nn_ops.py", line 4364, in _softmax_cross_entropy_with_logits
name=name)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/framework/op_def_library.py", line 787, in _apply_op_helper
op_def=op_def)
File "/root/tf-1.4/local/lib/python2.7/site-packages/tensorflow/python/framework/ops.py", line 2956, in create_op
op_def=op_def)
File "/root/tf-1.4/local/lib/python2.7/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): logits and labels must be same size: logits_size=[196,101] labels_size=[203,101]
[[Node: seq2seq/parallel_1/seq2seq/SoftmaxCrossEntropyWithLogits = SoftmaxCrossEntropyWithLogits[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"](seq2seq/parallel_1/seq2seq/Reshape, seq2seq/parallel_1/seq2seq/Reshape_1)]]

# 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
        )
    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)