Problem with running ANI simulations

Hi, I've opened an issue in https://github.com/openmm/openmm-torch/issues/26 about running simulations with ANI potentials but I thought I'd submit a ticket here since it might be more appropriate. I tried following the example in the docstring for creating a mixed system and run the code below.

prmtop = AmberPrmtopFile("molecule.prmtop")
inpcrd = AmberInpcrdFile("molecule.rst7")

potential = MLPotential('ani2x')
ml_atoms  = [atom.index for atom in prmtop.topology.atoms() if atom.residue.name == "MOL"]
mm_system = prmtop.createSystem(nonbondedMethod=NoCutoff)
ml_system = potential.createMixedSystem(prmtop.topology, mm_system, ml_atoms)

integrator = LangevinIntegrator(
    298.15 * unit.kelvin, 
    1 / unit.picosecond, 
    1.0 * unit.femtosecond,
)
simulation = Simulation(
    prmtop.topology,
    ml_system,
    integrator,
    Platform.getPlatformByName("CUDA"),
)
simulation.context.setPositions(inpcrd.positions)
simulation.minimizeEnergy()

But, I get the following error:

---------------------------------------------------------------------------
OpenMMException                           Traceback (most recent call last)
<ipython-input-6-e4ae2128ac25> in <module>
     27 )
     28 simulation.context.setPositions(inpcrd.positions)
---> 29 simulation.minimizeEnergy()

~/anaconda3/envs/openmm-torch/lib/python3.9/site-packages/simtk/openmm/app/simulation.py in minimizeEnergy(self, tolerance, maxIterations)
    126             to how many iterations it takes.
    127         """
--> 128         mm.LocalEnergyMinimizer.minimize(self.context, tolerance, maxIterations)
    129 
    130     def step(self, steps):

~/anaconda3/envs/openmm-torch/lib/python3.9/site-packages/simtk/openmm/openmm.py in minimize(context, tolerance, maxIterations)
  18990             the maximum number of iterations to perform. If this is 0, minimation is continued until the results converge without regard to how many iterations it takes. The default value is 0.
  18991         """
> 18992         return _openmm.LocalEnergyMinimizer_minimize(context, tolerance, maxIterations)
  18993     __swig_destroy__ = _openmm.delete_LocalEnergyMinimizer
  18994 

OpenMMException: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/torchani/aev.py", line 35, in forward
      _2, cell0 = False, unchecked_cast(Tensor, cell)
    if _2:
      aev0 = __torch__.torchani.aev.compute_aev(species, coordinates, self.triu_index, (self).constants(), (7, 16, 112, 32, 896), None, )
             ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
      aev = aev0
    else:
  File "code/__torch__/torchani/aev.py", line 110, in compute_aev
  species120 = torch.index_select(species12, 1, even_closer_indices)
  vec1 = torch.index_select(vec, 0, even_closer_indices)
  central_atom_index, pair_index12, sign12, = _7(atom_index122, )
                                              ~~ <--- HERE
  _22 = torch.slice(species120, 0, 0, 9223372036854775807, 1)
  _23 = annotate(List[Optional[Tensor]], [None, pair_index12])
  File "code/__torch__/torchani/aev.py", line 267, in triple_by_molecule
  sorted_local_index12 = torch.index(_85, _86)
  _87 = torch.index_select(_78(counts, ), 0, pair_indices)
  sorted_local_index120 = torch.add_(sorted_local_index12, _87, alpha=1)
                          ~~~~~~~~~~ <--- HERE
  _88 = annotate(List[Optional[Tensor]], [sorted_local_index120])
  local_index12 = torch.index(rev_indices, _88)

Traceback of TorchScript, original code (most recent call last):
  File "/home/jsetiadi/anaconda3/envs/openmm-torch/lib/python3.9/site-packages/torchani/aev.py", line 301, in compute_aev

    # compute angular aev
    central_atom_index, pair_index12, sign12 = triple_by_molecule(atom_index12)
                                               ~~~~~~~~~~~~~~~~~~ <--- HERE
    species12_small = species12[:, pair_index12]
    vec12 = vec.index_select(0, pair_index12.view(-1)).view(2, -1, 3) * sign12.unsqueeze(-1)
  File "/home/jsetiadi/anaconda3/envs/openmm-torch/lib/python3.9/site-packages/torchani/aev.py", line 246, in triple_by_molecule
    mask = (torch.arange(intra_pair_indices.shape[2], device=ai1.device) < pair_sizes.unsqueeze(1)).flatten()
    sorted_local_index12 = intra_pair_indices.flatten(1, 2)[:, mask]
    sorted_local_index12 += cumsum_from_zero(counts).index_select(0, pair_indices)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

    # unsort result from last part
RuntimeError: The size of tensor a (697) must match the size of tensor b (717) at non-singleton dimension 1

Is this a problem with how I set up my code above through the openmm-ml plugin or is this a bug in TorchANI?

molecule.zip

openmm / openmm-ml

Problem with running ANI simulations #5