I am hitting some error involving XLA and data.
This is a reproducible short test:
    data = tf.data.Dataset.from_tensor_slices(list(range(10)))
    data = data.repeat()
    iterator = iter(data)

    @tf.function(jit_compile=False)
    def run(it):
        t = 0
        for _ in range(10):
            x = next(it)
            t += x
        return t

    @tf.function(jit_compile=True)
    def run_compiled(it):
        t = 0
        for _ in range(10):
            x = next(it)
            t += x
        return t

    print(run(iterator))
    print(run_compiled(iterator))

Basically the non-compiled version (run) works fine, while the compiled one (run_compiled) is returning the following error:
tensorflow.python.framework.errors_impl.InternalError: Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to /job:localhost/replica:0/task:0/device:GPU:0 in order to run __inference_run_compiled_61: No unary variant device copy function found for direction: 1 and Variant type_index: tensorflow::ResourceDeleter [Op:__inference_run_compiled_61]

I am in the following tf version: 2.7.0-dev20210806
Any idea on what is going on?
Thanks!