StateBasedSentenceBreaker Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string

StateBasedSentenceBreaker's break_sentences fails when run on a GPU with Invalid argument: During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string. Which references a map_fn call here.

System information

• Have I written custom code (as opposed to using a stock example script
  provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 20.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue
  happens on a mobile device: n/a
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.7.1
• Python version: 3.8.10
• Bazel version (if compiling from source): n/a
• GCC/Compiler version (if compiling from source): n/a
• CUDA/cuDNN version: 11.2 / 8.1
• GPU model and memory: NVIDIA GeForce RTX 3090 - 24G & NVIDIA GeForce GTX 1080 8G
• Exact command to reproduce:

import tensorflow as tf
import tensorflow_text as text


@tf.function(input_signature=[tf.TensorSpec(shape=(None,), dtype=tf.string, name='doc')])
def split_sentences(doc):
    splitter = text.StateBasedSentenceBreaker()
    return splitter.break_sentences(doc)

doc = tf.constant(["Hello this is sentence 1. This is sentence 2."])
print(split_sentences(doc))

Other info/logs:
Traceback:

Traceback (most recent call last):
  File "<path>/.config/JetBrains/PyCharm2021.1/scratches/scratch_30.py", line 11, in <module>
    print(split_sentences(doc))
  File "<path>/venv/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "<path>/venv/lib/python3.8/site-packages/tensorflow/python/eager/execute.py", line 58, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) INVALID_ARGUMENT:  During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
	 [[node map/while/TensorArrayV2Write/TensorListSetItem
 (defined at <path>/venv/lib/python3.8/site-packages/tensorflow_text/python/ops/state_based_sentence_breaker_op.py:120)
]]
	 [[Assert_2/AssertGuard/pivot_f/_99/_355]]
  (1) INVALID_ARGUMENT:  During Variant Host->Device Copy: non-DMA-copy attempted of tensor type: string
	 [[node map/while/TensorArrayV2Write/TensorListSetItem
 (defined at <path>/venv/lib/python3.8/site-packages/tensorflow_text/python/ops/state_based_sentence_breaker_op.py:120)
]]
0 successful operations.
0 derived errors ignored. [Op:__inference_split_sentences_909]

Errors may have originated from an input operation.
Input Source operations connected to node map/while/TensorArrayV2Write/TensorListSetItem:
In[0] map/while/Placeholder_1:	
In[1] map/while/Placeholder:	
In[2] map/while/RaggedTensorToVariant:

Operation defined at: (most recent call last)
>>>   File "<path>/.config/JetBrains/PyCharm2021.1/scratches/scratch_30.py", line 11, in <module>
>>>     print(split_sentences(doc))
>>> 
>>>   File "<path>/.config/JetBrains/PyCharm2021.1/scratches/scratch_30.py", line 8, in split_sentences
>>>     return splitter.break_sentences(doc)
>>> 
>>>   File "<path>/venv/lib/python3.8/site-packages/tensorflow_text/python/ops/state_based_sentence_breaker_op.py", line 61, in break_sentences
>>>     results, _, _ = self.break_sentences_with_offsets(doc)
>>> 
>>>   File "<path>/venv/lib/python3.8/site-packages/tensorflow_text/python/ops/state_based_sentence_breaker_op.py", line 120, in break_sentences_with_offsets
>>>     fragment_text = map_fn.map_fn(
>>> 

Input Source operations connected to node map/while/TensorArrayV2Write/TensorListSetItem:
In[0] map/while/Placeholder_1:	
In[1] map/while/Placeholder:	
In[2] map/while/RaggedTensorToVariant:

Operation defined at: (most recent call last)
>>>   File "<path>/.config/JetBrains/PyCharm2021.1/scratches/scratch_30.py", line 11, in <module>
>>>     print(split_sentences(doc))
>>> 
>>>   File "<path>/.config/JetBrains/PyCharm2021.1/scratches/scratch_30.py", line 8, in split_sentences
>>>     return splitter.break_sentences(doc)
>>> 
>>>   File "<path>/venv/lib/python3.8/site-packages/tensorflow_text/python/ops/state_based_sentence_breaker_op.py", line 61, in break_sentences
>>>     results, _, _ = self.break_sentences_with_offsets(doc)
>>> 
>>>   File "<path>/venv/lib/python3.8/site-packages/tensorflow_text/python/ops/state_based_sentence_breaker_op.py", line 120, in break_sentences_with_offsets
>>>     fragment_text = map_fn.map_fn(
>>> 

Function call stack:
split_sentences -> map_while_body_853 -> split_sentences -> map_while_body_853

Start with tensorflow_text/python/ops/state_based_sentence_breaker_op.py, especially the break_sentences_with_offsets map_fn call around lines 120-125, and reproduce the issue using the provided tf.function and GPU configuration. Trace the map/while TensorArray error and verify that StateBasedSentenceBreaker.break_sentences works on a GPU without the InvalidArgumentError while returning the expected sentence output.