Speed Bottleneck due to dataset

[jacklu333333]

Hi, I am using TFUniformReplayBuffer with DynamicEpisodeDriverwith the former config device as GPU. However, in mytf.profile` it shows that the following code is actually running on the CPU.

dataset = replay_buffer.as_dataset(
    num_parallel_calls=16, 
    sample_batch_size=batch_size,
    num_steps=n_step_update + 1,
    ).prefetch(10)
iterator = iter(dataset)

experience, buffer_info = next(iterator)
train_loss = agent.train(experience)

Is there any way that I can make the iteration of the dataset run in the GPU instead of sending them back and forth?