Improving GPU usage for molecule optimization

[and-tos]

When optimizing a molecule, I observe GPU usages of around 10%. Any advice on how to improve this? I tried a dataloder, but get an error that atoms does not have a keys attribute.

My code looks like this:

def ccdc2ase(ccdc_mol):
    ccdc_atoms = ccdc_mol.atoms
    ase_atoms = ase.Atoms(
        positions=[a.coordinates for a in ccdc_atoms],
        numbers=[a.atomic_number for a in ccdc_atoms],
        charges=[a.formal_charge for a in ccdc_atoms],
    )
    ase_atoms.info["charge"] = ccdc_mol.formal_charge
    ase_atoms.info["spin"] = 1
    return ase_atoms

predictor = pretrained_mlip.load_predict_unit(
        "esen_md_direct_all.pt",
        device="cuda",
    )
calc = FAIRChemCalculator(predictor, task_name="omol")
atoms = ccdc2ase(ccdc_mol)
atoms.calc = calc
opt = FIRE(atoms)
try:
    opt.run(fmax=0.05, steps=100)  # steps=100
    energy = atoms.get_potential_energy() * 23.060548
except ValueError as e:
      print(f"Optimization did not converge: {e}")

[zulissimeta]

Please see the tutorial on how to use batched inference here:
https://fair-chem.github.io/core/common_tasks/batch_inference.html

Tagging @lbluque who probably has the most experience running batched relaxations in the past.

[mshuaibii]

This is the expected behavior and not an actual concern. Since inference is only ever happening on a single system at a time during an optimization, GPU utilization will be low.