How to ensemble multiple jit models?

I’m trying to deploy an ensemble model that consists of 5 base models. Sources like this and this provide great examples of building ensemble PyTorch models, but is there a way for us to either:

  • JIT (i.e., jit script and optimize_for_inference) an ensemble model (that is based on vmap); or
  • ensemble multiple JIT models?

I tried the former, which led to an error saying “Compiled functions can’t take variable number of arguments or use keyword-only arguments with defaults”. The latter approach led to an error saying attribute “” does not exist in these jit scripted models.

Any help is appreciated!