[Feature Request] Add option to use `functorch` for `BasicEnsemble`
luisenp opened this issue ยท 2 comments
luisenp commented
๐ Feature Request
Change BasicEnsemble
to optionally use functorch.vmap
.
Motivation and Pitch
Is your feature request related to a problem? Please describe.
BasicEnsemble
lets the user provide arbitrary models, which are stacked together using a very naive loop-based implementation. We should be able to do this more efficiently now using functorch
.
natolambert commented
Torch has a vmap
??? !!!