Is there a pytorch api like 'tf.map_fn' of tensorflow?

Is there a pytorch api like ‘tf.map_fn’ of tensorflow that I can do some duplicate operations parallelly on GPU?

For example, I have 64 tasks in one program, and each of the task have the same input data shape and same cnn network, but with different weights and biases, run these tasks sequencely is a easy way, but it is too slow,so I want to run the these tasks parallelly on GPU. In tensorflow, it can be realized by ‘tf.map_fn’, how can I do it in pytorch?

Thanks for any advice you have!

2 Likes