tf.raw_ops.ExperimentalDatasetToTFRecord: Aborted (core dumped)
Hi @nonenone12135 ,

Replicated the reported check fail and attached screenshot for reference.

<img width="1502" alt="Screenshot 2024-03-19 at 4 14 58 PM" src="https://github.com/tensorflow/tensorflow/assets/116063290/c412edb3-321a-4111-a30c-7d5d63de1492">

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

binary

### TensorFlow version

tf 2.15

### Custom code

Yes

### OS platform and distribution

Ubuntu 20.04

### Mobile device

_No response_

### Python version

3.9

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Under specific input, `tf.raw_ops.ExperimentalDatasetToTFRecord` encounters "Aborted (core dumped)".

### Standalone code to reproduce the issue

```shell
import tensorflow as tf

input_data = tf.data.Dataset.range(10).map(lambda x: tf.strings.as_string(x))
input_data = input_data.batch(2)  

filename = tf.constant("output.tfrecord")
compression_type = tf.constant("GZIP")
tf.raw_ops.ExperimentalDatasetToTFRecord(input_dataset=input_data._variant_tensor, filename=filename, compression_type=compression_type)
```


### Relevant log output

```shell
2024-03-14 05:47:55.415453: F tensorflow/core/framework/tensor.cc:852] Check failed: 1 == NumElements() (1 vs. 2)Must have a one element tensor
Aborted (core dumped)
```
