# Buggy code
def _lm_specific_token_cleanup(self, token_strings: List[Text]) -> List[Text]:
    from rasa.nlu.utils.hugging_face.registry import model_tokens_cleaners

    return model_tokens_cleaners[self.model_name](token_strings)


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
Traceback (most recent call last):
File "/home/ec2-user/anaconda3/envs/rasa/bin/rasa", line 11, in <module>
load_entry_point('rasa', 'console_scripts', 'rasa')()
File "/home/ec2-user/rasa/rasa/__main__.py", line 91, in main
cmdline_arguments.func(cmdline_arguments)
File "/home/ec2-user/rasa/rasa/cli/train.py", line 140, in train_nlu
persist_nlu_training_data=args.persist_nlu_data,
File "/home/ec2-user/rasa/rasa/train.py", line 414, in train_nlu
persist_nlu_training_data,
File "uvloop/loop.pyx", line 1456, in uvloop.loop.Loop.run_until_complete
File "/home/ec2-user/rasa/rasa/train.py", line 445, in _train_nlu_async
persist_nlu_training_data=persist_nlu_training_data,
File "/home/ec2-user/rasa/rasa/train.py", line 474, in _train_nlu_with_validated_data
persist_nlu_training_data=persist_nlu_training_data,
File "/home/ec2-user/rasa/rasa/nlu/train.py", line 86, in train
interpreter = trainer.train(training_data, **kwargs)
File "/home/ec2-user/rasa/rasa/nlu/model.py", line 191, in train
updates = component.train(working_data, self.config, **context)
File "/home/ec2-user/rasa/rasa/nlu/classifiers/diet_classifier.py", line 643, in train
self.component_config[BATCH_STRATEGY],
File "/home/ec2-user/rasa/rasa/utils/tensorflow/models.py", line 81, in fit
) = self._get_tf_train_functions(eager, model_data, batch_strategy)
File "/home/ec2-user/rasa/rasa/utils/tensorflow/models.py", line 241, in _get_tf_train_functions
train_dataset_function, self.train_on_batch, eager, "train"
File "/home/ec2-user/rasa/rasa/utils/tensorflow/models.py", line 223, in _get_tf_call_model_function
tf_call_model_function(next(iter(init_dataset)))
File "/home/ec2-user/anaconda3/envs/rasa/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py", line 568, in __call__
result = self._call(*args, **kwds)
File "/home/ec2-user/anaconda3/envs/rasa/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py", line 632, in _call
return self._stateless_fn(*args, **kwds)
File "/home/ec2-user/anaconda3/envs/rasa/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 2363, in __call__
return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
File "/home/ec2-user/anaconda3/envs/rasa/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1611, in _filtered_call
self.captured_inputs)
File "/home/ec2-user/anaconda3/envs/rasa/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1692, in _call_flat
ctx, args, cancellation_manager=cancellation_manager))
File "/home/ec2-user/anaconda3/envs/rasa/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 545, in call
ctx=ctx)
File "/home/ec2-user/anaconda3/envs/rasa/lib/python3.6/site-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:  Incompatible shapes: [1,6] vs. [1,7]
[[{{node cond/else/_83/add_1}}]]
[[crf/cond/else/_72/rnn/transpose_2/_260]]
(1) Invalid argument:  Incompatible shapes: [1,6] vs. [1,7]
[[{{node cond/else/_83/add_1}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_on_batch_36091]

Function call stack:
train_on_batch -> train_on_batch

# Fixed code
def _lm_specific_token_cleanup(
    self, split_token_ids: List[int], token_strings: List[Text]
) -> Tuple[List[int], List[Text]]:
    from rasa.nlu.utils.hugging_face.registry import model_tokens_cleaners

    return model_tokens_cleaners[self.model_name](split_token_ids, token_strings)