# Buggy code
def _get_docs_for_batch(
    self, batch_examples: List[Message], attribute: Text
) -> List[Dict[Text, Any]]:
    """Compute language model docs for all examples in the batch.

    Args:
        batch_examples: Batch of message objects for which language model docs
        need to be computed.
        attribute: Property of message to be processed, one of ``TEXT`` or
        ``RESPONSE``.

    Returns:
        List of language model docs for each message in batch.
    """

    batch_tokens, batch_token_ids = self._get_token_ids_for_batch(
        batch_examples, attribute
    )

    (
        batch_sentence_features,
        batch_sequence_features,
    ) = self._get_model_features_for_batch(batch_token_ids, batch_tokens)

    # A doc consists of
    # {'token_ids': ..., 'tokens': ..., 'sequence_features': ...,
    # 'sentence_features': ...}
    batch_docs = []
    for index in range(len(batch_examples)):
        doc = {
            TOKEN_IDS: batch_token_ids[index],
            TOKENS: batch_tokens[index],
            SEQUENCE_FEATURES: batch_sequence_features[index],
            SENTENCE_FEATURES: np.reshape(batch_sentence_features[index], (1, -1)),
        }
        batch_docs.append(doc)

    return batch_docs


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
2020-07-16 13:05:16 INFO     rasa.nlu.model  - Starting to train component HFTransformersNLP
Traceback (most recent call last):
File "/home/ubuntu/anaconda3/bin/rasa", line 8, in <module>
sys.exit(main())
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/__main__.py", line 92, in main
cmdline_arguments.func(cmdline_arguments)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/cli/train.py", line 140, in train_nlu
persist_nlu_training_data=args.persist_nlu_data,
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/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/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/train.py", line 453, in _train_nlu_async
persist_nlu_training_data=persist_nlu_training_data,
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/train.py", line 482, in _train_nlu_with_validated_data
persist_nlu_training_data=persist_nlu_training_data,
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/train.py", line 90, in train
interpreter = trainer.train(training_data, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/model.py", line 191, in train
updates = component.train(working_data, self.config, **context)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 472, in train
batch_docs = self._get_docs_for_batch(batch_messages, attribute)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 423, in _get_docs_for_batch
) = self._get_model_features_for_batch(batch_token_ids)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 387, in _get_model_features_for_batch
batch_attention_mask, padded_token_ids
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 347, in _compute_batch_sequence_features
np.array(padded_token_ids), attention_mask=np.array(batch_attention_mask)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
outputs = self.call(cast_inputs, *args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 708, in call
outputs = self.bert(inputs, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
outputs = self.call(cast_inputs, *args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 570, in call
embedding_output = self.embeddings([input_ids, position_ids, token_type_ids, inputs_embeds], training=training)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
outputs = self.call(cast_inputs, *args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 150, in call
return self._embedding(inputs, training=training)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 173, in _embedding
position_embeddings = self.position_embeddings(position_ids)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
outputs = self.call(cast_inputs, *args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/layers/embeddings.py", line 184, in call
out = embedding_ops.embedding_lookup(self.embeddings, inputs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/embedding_ops.py", line 323, in embedding_lookup
transform_fn=None)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/embedding_ops.py", line 137, in _embedding_lookup_and_transform
array_ops.gather(params[0], ids, name=name), ids, max_norm)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/util/dispatch.py", line 180, in wrapper
return target(*args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py", line 4106, in gather
return params.sparse_read(indices, name=name)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py", line 645, in sparse_read
self._handle, indices, dtype=self._dtype, name=name)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_resource_variable_ops.py", line 556, in resource_gather
_ops.raise_from_not_ok_status(e, name)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 6606, in raise_from_not_ok_status
six.raise_from(core._status_to_exception(e.code, message), None)
File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0,512] = 512 is not in [0, 512) [Op:ResourceGather] name: tf_bert_model/bert/embeddings/position_embeddings/embedding_lookup/

# Fixed code
def _get_docs_for_batch(
    self,
    batch_examples: List[Message],
    attribute: Text,
    inference_mode: bool = False,
) -> List[Dict[Text, Any]]:
    """Compute language model docs for all examples in the batch.

    Args:
        batch_examples: Batch of message objects for which language model docs
        need to be computed.
        attribute: Property of message to be processed, one of ``TEXT`` or
        ``RESPONSE``.
        inference_mode: Whether the call is during inference or during training.


    Returns:
        List of language model docs for each message in batch.
    """

    batch_tokens, batch_token_ids = self._get_token_ids_for_batch(
        batch_examples, attribute
    )

    (
        batch_sentence_features,
        batch_sequence_features,
    ) = self._get_model_features_for_batch(
        batch_token_ids, batch_tokens, batch_examples, attribute, inference_mode
    )

    # A doc consists of
    # {'token_ids': ..., 'tokens': ..., 'sequence_features': ...,
    # 'sentence_features': ...}
    batch_docs = []
    for index in range(len(batch_examples)):
        doc = {
            TOKEN_IDS: batch_token_ids[index],
            TOKENS: batch_tokens[index],
            SEQUENCE_FEATURES: batch_sequence_features[index],
            SENTENCE_FEATURES: np.reshape(batch_sentence_features[index], (1, -1)),
        }
        batch_docs.append(doc)

    return batch_docs