Hi,
If the latency doesn’t matter, I think the snippet below will work.
weights = self.weight_network(x)
f = sum(w * m(x) for w, m in zip(weights, self.base_modules))

I tried this here: https://colab.research.google.com/drive/163Id8tsSofW4c_Mm6_7X_XThJadTpn6e?usp=sharing 3