Repository: OpenNMT/OpenNMT-tf

Issue:
InvalidArgumentError when running evaluation with noise decoding
I use OpenNMT-tf to train a transformer model. The traning process goes well, but en error occurs when running evaluation.

Here is the log

```bash
INFO:tensorflow:Saved checkpoint /opt/algo_nfs/kdd_luozhouyang/jdrewrite/model/transformer/ckpt-0
WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/resource_variable_ops.py:1781: calResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removuture version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
2019-10-16 02:56:01.463522: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:143] Filling up shuffle buffer (this may tae): 1521766 of 1799870
2019-10-16 02:56:03.298341: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:193] Shuffle buffer filled.
INFO:tensorflow:Step = 100 ; source words/s = 10616, target words/s = 289 ; Learning rate = 0.000100 ; Loss = 5.583332
INFO:tensorflow:Step = 200 ; source words/s = 77108, target words/s = 1940 ; Learning rate = 0.000100 ; Loss = 5.666554
INFO:tensorflow:Step = 300 ; source words/s = 78702, target words/s = 1986 ; Learning rate = 0.000100 ; Loss = 5.430358
INFO:tensorflow:Step = 400 ; source words/s = 77159, target words/s = 1919 ; Learning rate = 0.000100 ; Loss = 5.042500
INFO:tensorflow:Step = 500 ; source words/s = 76053, target words/s = 2014 ; Learning rate = 0.000100 ; Loss = 4.841046
INFO:tensorflow:Step = 600 ; source words/s = 78079, target words/s = 1891 ; Learning rate = 0.000100 ; Loss = 4.983972
INFO:tensorflow:Step = 700 ; source words/s = 77481, target words/s = 2021 ; Learning rate = 0.000100 ; Loss = 4.961732
INFO:tensorflow:Step = 800 ; source words/s = 77352, target words/s = 1915 ; Learning rate = 0.000100 ; Loss = 4.508535
INFO:tensorflow:Step = 900 ; source words/s = 76161, target words/s = 1987 ; Learning rate = 0.000100 ; Loss = 3.886328
INFO:tensorflow:Step = 1000 ; source words/s = 76549, target words/s = 1901 ; Learning rate = 0.000100 ; Loss = 4.771955
INFO:tensorflow:Step = 1100 ; source words/s = 79043, target words/s = 1934 ; Learning rate = 0.000100 ; Loss = 3.687625
INFO:tensorflow:Step = 1200 ; source words/s = 78311, target words/s = 1921 ; Learning rate = 0.000100 ; Loss = 3.797333
INFO:tensorflow:Step = 1300 ; source words/s = 78706, target words/s = 1999 ; Learning rate = 0.000100 ; Loss = 3.763635
INFO:tensorflow:Step = 1400 ; source words/s = 78690, target words/s = 2004 ; Learning rate = 0.000100 ; Loss = 3.591205
INFO:tensorflow:Step = 1500 ; source words/s = 79119, target words/s = 1951 ; Learning rate = 0.000100 ; Loss = 3.379777
INFO:tensorflow:Step = 1600 ; source words/s = 77603, target words/s = 1934 ; Learning rate = 0.000100 ; Loss = 3.521855
INFO:tensorflow:Step = 1700 ; source words/s = 76440, target words/s = 2002 ; Learning rate = 0.000100 ; Loss = 2.911270
INFO:tensorflow:Step = 1800 ; source words/s = 77462, target words/s = 1838 ; Learning rate = 0.000100 ; Loss = 3.304400
INFO:tensorflow:Step = 1900 ; source words/s = 77666, target words/s = 1945 ; Learning rate = 0.000100 ; Loss = 3.305823
INFO:tensorflow:Step = 2000 ; source words/s = 76761, target words/s = 1893 ; Learning rate = 0.000100 ; Loss = 2.937279
INFO:tensorflow:Step = 2100 ; source words/s = 77877, target words/s = 2019 ; Learning rate = 0.000100 ; Loss = 2.932232
INFO:tensorflow:Step = 2200 ; source words/s = 76121, target words/s = 1924 ; Learning rate = 0.000100 ; Loss = 3.219757
INFO:tensorflow:Step = 2300 ; source words/s = 77865, target words/s = 1909 ; Learning rate = 0.000100 ; Loss = 2.884820
INFO:tensorflow:Step = 2400 ; source words/s = 75515, target words/s = 1970 ; Learning rate = 0.000100 ; Loss = 2.969361
INFO:tensorflow:Step = 2500 ; source words/s = 77161, target words/s = 1861 ; Learning rate = 0.000100 ; Loss = 2.742362
INFO:tensorflow:Step = 2600 ; source words/s = 76826, target words/s = 1941 ; Learning rate = 0.000100 ; Loss = 2.909730
INFO:tensorflow:Step = 2700 ; source words/s = 76092, target words/s = 1955 ; Learning rate = 0.000100 ; Loss = 3.543749
INFO:tensorflow:Step = 2800 ; source words/s = 77019, target words/s = 1937 ; Learning rate = 0.000100 ; Loss = 2.521420
INFO:tensorflow:Step = 2900 ; source words/s = 77562, target words/s = 1834 ; Learning rate = 0.000100 ; Loss = 2.453693
INFO:tensorflow:Step = 3000 ; source words/s = 76726, target words/s = 1950 ; Learning rate = 0.000100 ; Loss = 2.646461
INFO:tensorflow:Step = 3100 ; source words/s = 76233, target words/s = 1901 ; Learning rate = 0.000100 ; Loss = 3.034106
INFO:tensorflow:Step = 3200 ; source words/s = 75719, target words/s = 1955 ; Learning rate = 0.000100 ; Loss = 3.000289
INFO:tensorflow:Step = 3300 ; source words/s = 76795, target words/s = 1940 ; Learning rate = 0.000100 ; Loss = 2.223612
INFO:tensorflow:Step = 3400 ; source words/s = 77994, target words/s = 1928 ; Learning rate = 0.000100 ; Loss = 2.867700
INFO:tensorflow:Step = 3500 ; source words/s = 75764, target words/s = 1936 ; Learning rate = 0.000100 ; Loss = 2.304001
INFO:tensorflow:Step = 3600 ; source words/s = 78271, target words/s = 1863 ; Learning rate = 0.000100 ; Loss = 2.193941
INFO:tensorflow:Step = 3700 ; source words/s = 76626, target words/s = 2013 ; Learning rate = 0.000100 ; Loss = 2.812444
INFO:tensorflow:Step = 3800 ; source words/s = 77158, target words/s = 1861 ; Learning rate = 0.000100 ; Loss = 2.488263
INFO:tensorflow:Step = 3900 ; source words/s = 77520, target words/s = 1848 ; Learning rate = 0.000100 ; Loss = 2.229686
INFO:tensorflow:Step = 4000 ; source words/s = 76969, target words/s = 1936 ; Learning rate = 0.000100 ; Loss = 3.432198
INFO:tensorflow:Step = 4100 ; source words/s = 77902, target words/s = 1921 ; Learning rate = 0.000100 ; Loss = 2.815073
INFO:tensorflow:Step = 4200 ; source words/s = 78302, target words/s = 1989 ; Learning rate = 0.000100 ; Loss = 2.300324
INFO:tensorflow:Step = 4300 ; source words/s = 78127, target words/s = 2023 ; Learning rate = 0.000100 ; Loss = 2.335404
INFO:tensorflow:Step = 4400 ; source words/s = 79337, target words/s = 1939 ; Learning rate = 0.000100 ; Loss = 2.305170
INFO:tensorflow:Step = 4500 ; source words/s = 77957, target words/s = 1956 ; Learning rate = 0.000100 ; Loss = 2.216660
INFO:tensorflow:Step = 4600 ; source words/s = 77612, target words/s = 1936 ; Learning rate = 0.000100 ; Loss = 2.549562
INFO:tensorflow:Step = 4700 ; source words/s = 77886, target words/s = 1901 ; Learning rate = 0.000100 ; Loss = 2.596675
INFO:tensorflow:Step = 4800 ; source words/s = 77120, target words/s = 1955 ; Learning rate = 0.000100 ; Loss = 1.904201
INFO:tensorflow:Step = 4900 ; source words/s = 76252, target words/s = 1977 ; Learning rate = 0.000100 ; Loss = 2.229644
INFO:tensorflow:Step = 5000 ; source words/s = 78700, target words/s = 1816 ; Learning rate = 0.000100 ; Loss = 2.202982
INFO:tensorflow:Running evaluation for step 5000
2019-10-16 03:05:49.514388: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAlid argument: 2 root error(s) found.
  (0) Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
  (1) Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
0 successful operations.
0 derived errors ignored.
	 [[{{node transformer/map/TensorArrayUnstack/TensorListFromTensor/_96}}]]
	 [[transformer/map/while/body/_397/RaggedFromTensor/boolean_mask/concat/_252]]
2019-10-16 03:05:49.514983: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAlid argument: 2 root error(s) found.
  (0) Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
  (1) Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
0 successful operations.
0 derived errors ignored.
	 [[{{node transformer/map/TensorArrayUnstack/TensorListFromTensor/_96}}]]
Traceback (most recent call last):
  File "/usr/local/bin/onmt-main", line 10, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.6/dist-packages/opennmt/bin/main.py", line 189, in main
    checkpoint_path=args.checkpoint_path)
  File "/usr/local/lib/python3.6/dist-packages/opennmt/runner.py", line 205, in train
    export_on_best=eval_config.get("export_on_best"))
  File "/usr/local/lib/python3.6/dist-packages/opennmt/training.py", line 168, in __call__
    self._evaluate(evaluator, step, export_on_best=export_on_best)
  File "/usr/local/lib/python3.6/dist-packages/opennmt/training.py", line 180, in _evaluate
    metrics = evaluator(step)
  File "/usr/local/lib/python3.6/dist-packages/opennmt/evaluation.py", line 241, in __call__
    for loss, predictions, target in self._eval():  # pylint: disable=no-value-for-parameter
  File "/usr/local/lib/python3.6/dist-packages/opennmt/data/dataset.py", line 433, in _fun
    outputs = _tf_fun()
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/def_function.py", line 457, in __call__
    result = self._call(*args, **kwds)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/def_function.py", line 526, in _call
    return self._concrete_stateful_fn._filtered_call(canon_args, canon_kwds)  # pylint: disable=protected-access
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/function.py", line 1141, in _filtered_call
    self.captured_inputs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/function.py", line 1224, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/function.py", line 511, in call
    ctx=ctx)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/execute.py", line 67, in quick_execute
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  2 root error(s) found.
  (0) Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
  (1) Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
0 successful operations.
0 derived errors ignored.
	 [[{{node transformer/map/TensorArrayUnstack/TensorListFromTensor/_96}}]]
	 [[transformer/map/while/body/_397/RaggedFromTensor/boolean_mask/concat/_252]]
  (1) Invalid argument:  2 root error(s) found.
  (0) Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
  (1) Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
0 successful operations.
0 derived errors ignored.
	 [[{{node transformer/map/TensorArrayUnstack/TensorListFromTensor/_96}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference__tf_fun_60338]

Function call stack:
_tf_fun -> _tf_fun
```



Discussion:
Here is a issue of tensorflow [tf.function failed with tf.image.summary](https://github.com/tensorflow/tensorflow/issues/28007)
Can you post your configuration file?
Here is my config file:

```yml
model_dir: /opt/algo_nfs/kdd_luozhouyang/model/transformer

data:
  # (required for train run type).
  train_features_file: /opt/algo_nfs/kdd_luozhouyang/onmt/src.1800000.lower.char.train
  train_labels_file: /opt/algo_nfs/kdd_luozhouyang/onmt/tgt.1800000.lower.char.train

  # (optional) Pharaoh alignments of the training files.
  train_alignment:

  # (required for train_end_eval and eval run types).
  eval_features_file: /opt/algo_nfs/kdd_luozhouyang/onmt/src.100000.lower.char.eval
  eval_labels_file: /opt/algo_nfs/kdd_luozhouyang/onmt/tgt.100000.lower.char.eval

  # (optional) Models may require additional resource files (e.g. vocabularies).
  source_vocabulary: /opt/algo_nfs/kdd_luozhouyang/onmt/vocab.src.txt
  target_vocabulary: /opt/algo_nfs/kdd_luozhouyang/onmt/vocab.tgt.txt

  # (optional) Tokenization configuration (or path to a configuration file).
  # See also: https://github.com/OpenNMT/Tokenizer/blob/master/docs/options.md
  source_tokenization:
    type: SpaceTokenizer
    #params:
    #mode: space
      #joiner_annotate: true
      #segment_numbers: true
      #segment_alphabet_change: true

  target_tokenization:
    type: SpaceTokenizer
    #params:
    #mode: space
      #joiner_annotate: true
      #segment_numbers: true
      #segment_alphabet_change: true


  # (optional) Pretrained embedding configuration.
  #source_embedding:
  #  path: data/glove/glove-100000.txt
  #  with_header: True
  #  case_insensitive: True
  #  trainable: False

  # (optional) For sequence tagging tasks, the tagging scheme that is used (e.g. BIOES).
  # For supported schemes, additional evaluation metrics could be computed such as
  # precision, recall, etc. (accepted values: bioes; default: null).
  #tagging_scheme: bioes

# Model and optimization parameters.
params:
  # The optimizer class name in tf.keras.optimizers or tfa.optimizers.
  optimizer: Adam
  # (optional) Additional optimizer parameters as defined in their documentation.
  # If weight_decay is set, the optimizer will be extended with decoupled weight decay.
  optimizer_params:
    beta_1: 0.8
    beta_2: 0.998
  learning_rate: 1.0

  # (optional) If set, overrides all dropout values configured in the model definition.
  dropout: 0.3

  # (optional) List of layer to not optimize.
  #freeze_layers:
  #  - "encoder/layers/0"
  #  - "decoder/output_layer"

  # (optional) Weights regularization penalty (default: null).
  regularization:
    type: l2  # can be "l1", "l2", "l1_l2" (case-insensitive).
    scale: 1e-4  # if using "l1_l2" regularization, this should be a YAML list.

  # (optional) Average loss in the time dimension in addition to the batch dimension (default: False).
  average_loss_in_time: false

  # (optional) The type of learning rate decay (default: null). See:
  #  * https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules
  #  * opennmt/schedules/lr_schedules.py
  # This value may change the semantics of other decay options. See the documentation or the code.
  decay_type: NoamDecay
  # (optional unless decay_type is set) Decay parameters.
  decay_params:
    model_dim: 512
    warmup_steps: 4000
  # (optional) After how many steps to start the decay (default: 0).
  start_decay_steps: 50000

  # (optional) The learning rate minimum value (default: 0).
  minimum_learning_rate: 0.0001

  # (optional) Type of scheduled sampling (can be "constant", "linear", "exponential",
  # or "inverse_sigmoid", default: "constant").
  #scheduled_sampling_type: constant
  # (optional) Probability to read directly from the inputs instead of sampling categorically
  # from the output ids (default: 1).
  #scheduled_sampling_read_probability: 1
  # (optional unless scheduled_sampling_type is set) The constant k of the schedule.
  #scheduled_sampling_k: 0

  # (optional) The label smoothing value.
  label_smoothing: 0.1

  # (optional) Width of the beam search (default: 1).
  beam_width: 5
  # (optional) Number of hypotheses to return (default: 1). Set 0 to return all
  # available hypotheses. This value is also set by infer/n_best.
  num_hypotheses: 1
  # (optional) Length penaly weight to use during beam search (default: 0).
  length_penalty: 0.2
  # (optional) Coverage penaly weight to use during beam search (default: 0).
  coverage_penalty: 0.2
  # (optional) Sample predictions from the top K most likely tokens (requires
  # beam_width to 1). If 0, sample from the full output distribution (default: 1).
  sampling_topk: 1
  # (optional) High temperatures generate more random samples (default: 1).
  sampling_temperature: 1
  # (optional) Sequence of noise to apply to the decoding output. Each element
  # should be a noise type (can be: "dropout", "replacement", "permutation") and
  # the module arguments
  # (see http://opennmt.net/OpenNMT-tf/package/opennmt.data.noise.html)
  decoding_noise:
    - dropout: 0.1
    - replacement: [0.1, ｟unk｠]
    - permutation: 3
  # (optional) Define the subword marker. This is useful to apply noise at the
  # word level instead of the subword level (default: ￭).
  decoding_subword_token: ￭
  # (optional) Whether decoding_subword_token is used as a spacer (as in SentencePiece) or a joiner (as in BPE).
  # If unspecified, will infer  directly from decoding_subword_token.
  decoding_subword_token_is_spacer: false
  # (optional) Minimum length of decoded sequences, end token excluded (default: 0).
  minimum_decoding_length: 0
  # (optional) Maximum length of decoded sequences, end token excluded (default: 250).
  maximum_decoding_length: 30

  # (optional) Replace unknown target tokens by the original source token with the
  # highest attention (default: false).
  replace_unknown_target: false

  # (optional) The type of guided alignment cost to compute (can be: "null", "ce", "mse",
  # default: "null").
  guided_alignment_type: null
  # (optional) The weight of the guided alignment cost (default: 1).
  guided_alignment_weight: 1

  # (optional) Enable contrastive learning mode, see
  # https://www.aclweb.org/anthology/P19-1623 (default: false).
  # See also "decoding_subword_token" that is used by this mode.
  contrastive_learning: false
  # (optional) The value of the parameter eta in the max-margin loss (default: 0.1).
  max_margin_eta: 0.1

# Training options.
train:
  # (optional when batch_type=tokens) If not set, the training will search the largest
  # possible batch size.
  batch_size: 32
  # (optional) Batch size is the number of "examples" or "tokens" (default: "examples").
  batch_type: examples
  # (optional) Tune gradient accumulation to train with at least this effective batch size
  # (default: null).
  effective_batch_size: null

  # (optional) Save a checkpoint every this many steps (default: 5000).
  save_checkpoints_steps: 10000
  # (optional) How many checkpoints to keep on disk.
  keep_checkpoint_max: 5

  # (optional) Dump summaries and logs every this many steps (default: 100).
  save_summary_steps: 100

  # (optional) Maximum training step. If not set, train forever.
  max_step: 1000000
  # (optional) If true, makes a single pass over the training data (default: false).
  single_pass: false

  # (optional) The maximum length of feature sequences during training (default: null).
  maximum_features_length: 1500
  # (optional) The maximum length of label sequences during training (default: null).
  maximum_labels_length: 30

  # (optional) The width of the length buckets to select batch candidates from.
  # A smaller value means less padding and increased efficiency. (default: 1).
  length_bucket_width: 100

  # (optional) The number of elements from which to sample during shuffling (default: 500000).
  # Set 0 or null to disable shuffling, -1 to match the number of training examples.
  sample_buffer_size: -1

  # (optional) Number of checkpoints to average at the end of the training to the directory
  # model_dir/avg (default: 0).
  average_last_checkpoints: 6

# (optional) Evaluation options.
eval:
  # (optional) The batch size to use (default: 32).
  batch_size: 30

  # (optional) Evaluate every this many steps (default: 5000).
  steps: 5000

  # (optional) Save evaluation predictions in model_dir/eval/.
  save_eval_predictions: true
  # (optional) Evalutator or list of evaluators that are called on the saved evaluation predictions.
  # Available evaluators: bleu, rouge
  external_evaluators: bleu

  # (optional) Export a SavedModel when a metric has the best value so far (default: null).
  export_on_best: bleu

  # (optional) Early stopping condition.
  # Should be read as: stop the training if "metric" did not improve more
  # than "min_improvement" in the last "steps" evaluations.
  early_stopping:
    # (optional) The target metric name (default: "loss").
    metric: bleu
    # (optional) The metric should improve at least by this much to be considered as an improvement (default: 0)
    min_improvement: 0.01
    steps: 10

# (optional) Inference options.
infer:
  # (optional) The batch size to use (default: 1).
  batch_size: 10

  # (optional) For compatible models, the number of hypotheses to output (default: 1).
  # This sets the parameter params/num_hypotheses.
  n_best: 1
  # (optional) For compatible models, also output the score (default: false).
  with_scores: false
  # (optional) For compatible models, also output the alignments (can be: null, hard, soft,
  # default: null).
  with_alignments: null

  # (optional) The width of the length buckets to select batch candidates from.
  # If set, the test data will be sorted by length to increase the translation
  # efficiency. The predictions will still be outputted in order as they are
  # available (default: 0).
  length_bucket_width: 100

# (optional) Scoring options.
score:
  # (optional) The batch size to use (default: 64).
  batch_size: 64
  # (optional) Also report token-level cross entropy.
  with_token_level: false
  # (optional) Also output the alignments (can be: null, hard, soft, default: null).
  with_alignments: null
```
The error appears when using:

```yaml
params:
  decoding_noise:
    - dropout: 0.1
    - replacement: [0.1, ｟unk｠]
    - permutation: 3
```

Do you want to decode with noise? If not, remove this block and it should run fine.

On a related note, it looks like you copied all parameters from http://opennmt.net/OpenNMT-tf/configuration.html. I would not recommend doing that. Instead, start from an empty configuration file and progressively add the parameters you want to set.

Resolving patch:
diff --git a/opennmt/data/noise.py b/opennmt/data/noise.py
--- a/opennmt/data/noise.py
+++ b/opennmt/data/noise.py
@@ -54,6 +54,10 @@ def __call__(self, tokens, sequence_length=None, keep_shape=False):
     Returns:
       A tuple with the noisy version of :obj:`tokens` and the new lengths.
     """
+    with tf.device("cpu:0"):
+      return self._call(tokens, sequence_length, keep_shape)
+
+  def _call(self, tokens, sequence_length, keep_shape):
     rank = tokens.shape.ndims
     if rank == 1:
       input_length = tf.shape(tokens)[0]
@@ -77,7 +81,7 @@ def __call__(self, tokens, sequence_length=None, keep_shape=False):
       if sequence_length is None:
         raise ValueError("sequence_length must be passed for 2D inputs")
       tokens, sequence_length = tf.map_fn(
-          lambda arg: self(*arg, keep_shape=True),
+          lambda arg: self._call(*arg, keep_shape=True),
           (tokens, sequence_length),
           back_prop=False)
       if not keep_shape:
@@ -89,7 +93,7 @@ def __call__(self, tokens, sequence_length=None, keep_shape=False):
       original_shape = misc.shape_list(tokens)
       tokens = tf.reshape(tokens, [-1, original_shape[-1]])
       sequence_length = tf.reshape(sequence_length, [-1])
-      tokens, sequence_length = self(tokens, sequence_length, keep_shape=keep_shape)
+      tokens, sequence_length = self._call(tokens, sequence_length, keep_shape=keep_shape)
       tokens = tf.reshape(tokens, original_shape[:-1] + [-1])
       sequence_length = tf.reshape(sequence_length, original_shape[:-1])
       return tokens, sequence_length
diff --git a/opennmt/tokenizers/tokenizer.py b/opennmt/tokenizers/tokenizer.py
--- a/opennmt/tokenizers/tokenizer.py
+++ b/opennmt/tokenizers/tokenizer.py
@@ -76,6 +76,10 @@ def tokenize(self, text):
     Raises:
       ValueError: if the rank of :obj:`text` is greater than 1.
     """
+    with tf.device("cpu:0"):
+      return self._tokenize(text)
+
+  def _tokenize(self, text):
     if tf.is_tensor(text):
       rank = len(text.shape)
       if rank == 0:
@@ -112,6 +116,10 @@ def detokenize(self, tokens, sequence_length=None):
       ValueError: if :obj:`tokens` is a 2-D dense ``tf.Tensor`` and
         :obj:`sequence_length` is not set.
     """
+    with tf.device("cpu:0"):
+      return self._detokenize(tokens, sequence_length)
+
+  def _detokenize(self, tokens, sequence_length):
     if tf.is_tensor(tokens):
       rank = len(tokens.shape)
       if rank == 1:


Test patch:
diff --git a/opennmt/tests/noise_test.py b/opennmt/tests/noise_test.py
--- a/opennmt/tests/noise_test.py
+++ b/opennmt/tests/noise_test.py
@@ -60,14 +60,21 @@ def testWordPermutation(self, k):
       for i, v in enumerate(y.tolist()):
         self.assertLess(abs(int(v) - i), k)
 
-  def testWordNoising(self):
-    tokens = tf.constant([["a￭", "b", "c￭", "d", "￭e"], ["a", "b", "c", "", ""]])
-    lengths = tf.constant([5, 3])
+  @parameterized.expand([
+      [True, [["a￭", "b", "c￭", "d", "￭e"], ["a", "b", "c", "", ""]], [5, 3]],
+      [False, [["a￭", "b", "c￭", "d", "￭e"], ["a", "b", "c", "", ""]], [5, 3]],
+      [False, ["a￭", "b", "c￭", "d", "￭e"], None]
+  ])
+  def testWordNoising(self, as_function, tokens, lengths):
+    tokens = tf.constant(tokens)
+    if lengths is not None:
+      lengths = tf.constant(lengths, dtype=tf.int32)
     noiser = noise.WordNoiser()
     noiser.add(noise.WordDropout(0.1))
     noiser.add(noise.WordReplacement(0.1))
     noiser.add(noise.WordPermutation(3))
-    noisy_tokens, noisy_lengths = noiser(tokens, sequence_length=lengths, keep_shape=True)
+    noiser_fn = tf.function(noiser) if as_function else noiser
+    noisy_tokens, noisy_lengths = noiser_fn(tokens, sequence_length=lengths, keep_shape=True)
     tokens, noisy_tokens = self.evaluate([tokens, noisy_tokens])
     self.assertAllEqual(noisy_tokens.shape, tokens.shape)