Running multiple modules, as single module

Hi, I have a ModuleDictionary where each module take similar size input, x.
I have two use cases for it:

  1. For a single input x run all module in dictionary and create tensor by stacking all outputs.
  2. Give a ordered map of {key, x} run each x to module in dictionary for that key and stack the outputs in that order. I can change input of this case to be batch of x and corresponding list of key.

Currently I am doing this by looping on dictionary. Is there a better way to do this.

I think iterating the dict would be the right approach. If you are concerned about the performance of this loop, you could use a dict comprehension, which could yield a speedup, but I would recommend to profile the code first and check, if this loop is indeed the bottleneck in your code.

1 Like