`tf.data.Dataset` only supports Python-style iteration in eager mode or within tf.function.
@lmx666-gif,
I tried to execute the mentioned code and it was failing with the different error. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/5f35a72261875affb691a5549eaa5ca3/untitled1250.ipynb) and provide the complete code, dependencies and the tensorflow version you are using which helps to analyse the issue in an effective way. 

**tfds** is compatible with both eager and graph modes, but that doesn't mean they can be used the same way in either context.
 If you want a tensor representation of the inputs, you can use **outputs = dataset.make_one_shot_iterator().get_next().**
Also when you apply the `tf.function` decorator to a Python function that uses `tf.data.Dataset`, Tensorflow 2.0 will automatically convert the Python function into a graph that can be optimized for performance. This can lead to significant performance improvements, especially when working with large datasets.
Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. Closing this issue as stale. Please reopen if this is still a valid request. Thank you!
The first is the map function:
    def map_function(example):

        feature_map = {"wav_raw": tf.io.FixedLenFeature([], tf.string)}
        parsed_example = tf.io.parse_single_example(example, features=feature_map)

        wav_slice = tf.io.decode_raw(parsed_example["wav_raw"], out_type=tf.float64)
        wav_slice = tf.cast(wav_slice, tf.float32) / 2 ** 15

        return wav_slice

The second one is the training process:

     for epoch in range(args.num_epochs):
              
        trainset = tf.data.TFRecordDataset(args.trainset_tfrecords_path)
        trainset = trainset.map(map_func=map_function,
                                num_parallel_calls=num_cpus)  # num_parallel_calls should be number of cpu cores

        #trainset = trainset.shuffle(buffer_size=args.batch_size * 200, reshuffle_each_iteration=True)
        trainset = trainset.batch(batch_size=args.batch_size)
        trainset = trainset.prefetch(buffer_size=args.batch_size)


        # train_loss for each epoch
        train_loss_epoch = []
        train_loss = 0.0

        # record the train time for each epoch
        start = time.time()
        # MASK参数

        #EMA_MODEL来选择mask的index

        binary_mask = RandomMaskingGenerator(input_size,frame_length,mask_ratio)
        # bmr:0为掩码，1为未掩码；
        # bm_T:1为掩码,0为未掩码；
        # bm,bm_T = binary_mask()
        

        for step, _input in enumerate(trainset):
            bm, bm_T, _ = binary_mask.random_mask(_input.shape[0],alpha_e_max)#.totally_random_mask(_input.shape[0])
            print("_input",_input)
            # print("bm_shape", bm.shape)
            # print("bm_T_shape", bm_T.shape)
            loss_value = train_step(_input,_input*bm,bm_T)
            loss_float = float(loss_value)

            train_loss_epoch.append(loss_float)

            # Calculate the accumulated train loss value
            train_loss += loss_float



        # average train loss for each epoch
        train_loss /= (step + 1)
        train_loss_all.append(train_loss)

        # print log
        log = "train epoch {}/{}, train_loss = {:.06f}, time = {:.06f}"
The third one is the train_step function:

    @tf.function
    def train_step(_input,_input_mask,bm_T):
        with tf.GradientTape() as tape:
            enc_output, batch_mean, batch_var = sem_enc(_input_mask)
            #输入进去semantic decoder
            print("main:",enc_output)
            _output = sem_dec([enc_output, batch_mean, batch_var])
            loss_value = mse_loss(tf.multiply(_input, bm_T), _output)
            tf.print(loss_value)
            loss_whole = loss_value

        grads = tape.gradient(loss_whole, weights_all)  # compute gradients
        optimizer.apply_gradients(zip(grads, weights_all))  # update parameters

        return loss_whole



The _ouput generated by sem_dec() is the value I wanted change it from kerasTensor to TF.tensor; I could send you the whole codes, if you need, this is my email: limingxiao@link.cuhk.edu.cn. Thank you for your reply!

_Originally posted by @lmx666-gif in https://github.com/tensorflow/tensorflow/issues/61307#issuecomment-1643164557_
            