Your tensorflow dataset has a range method where you can specify the start, stop and step,
your pytorch dataset, requires you to give a chunk of data to each worker. and you can identify each worker using workerinfo if I’m not mistaken!
consider the following example :
import tensorflow as tf 
dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
dataset = dataset.map(lambda x: x*2)
worker1_data =dataset.range(0, 3, 1)
worker2_data = dataset.range(4, 8, 1)
...

and based on each worker you can use one of them.
you may also have a look at tf.data.Dataset.shard() method which was created for this very reason.