Any way to subset Ensemble to 2 or 4 of the 8 models?

[JSLJ23]

Hi TorchANI devs,

I was wondering if there was a way to subset the ensemble of 8 models to choose either 2 or 4? Using all 8 models seem to be causing GPU OOM errors while I understand that there is a way to just use a single model with the model_index parameter. Would there be any way to achieve a middle ground for this and use a subset of the 8 models so that less VRAM is required, yet I can still have the benefit of an Ensemble?

Thank you and I hope to hear from you all soon.

[yueyericardo]

Hi Joshua,

To select the first 2 models within the ensemble, it could be achived by

import torchani
model = torchani.models.ANI2x()
model.neural_networks = model.neural_networks[:2]

[JSLJ23]

Ok cool, I manage to find way to subset given a list of model indexes as well. Thank you for the prompt reply:)

model_indexes = [0,1,2,3]

models = []

for model_index in model_indexes:
    if 0 <= model_index <= 7:
        models.append(model.neural_networks[model_index])

model.neural_networks = Ensemble(models)