torch.load fails for AIMNet2 models (.jpt) with PyTorch >= 2.6 due to weights_only=True default

[EidRashed]

I got this error when I tried to use this fork https://github.com/peastman/openmm-ml/blob/aimnet2/openmmml/models/aimnet2potential.py

Traceback (most recent call last): File "/app/someproject/production/FC/test.py", line 607, in <module> Minimize() File "/app/someproject/production/FC/test.py", line 245, in Minimize system = create_mixed_system( File "/app/someproject/production/FC/test.py", line 226, in create_mixed_system hybrid_system = potential.createMixedSystem( File "/root/miniconda3/envs/workflow2/lib/python3.10/site-packages/openmmml/mlpotential.py", line 269, in createMixedSystem self._impl.addForces(topology, newSystem, atomList, forceGroup, **args) File "/root/miniconda3/envs/workflow2/lib/python3.10/site-packages/openmmml/models/aimnet2potential.py", line 62, in addForces model = torch.load('aimnet2_wb97m-d3_ens.jpt') File "/root/miniconda3/envs/workflow2/lib/python3.10/site-packages/torch/serialization.py", line 1495, in load raise RuntimeError( RuntimeError: Cannot use ``weights_only=True`` with TorchScript archives passed to ``torch.load``. In PyTorch 2.6, we changed the default value of the weights_onlyargument intorch.loadfromFalsetoTrue. Re-running torch.loadwithweights_onlyset toFalse will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.

[peastman]

Can you explain exactly what you're doing? If possible, provide a complete test case to reproduce it.

[EidRashed]

I'm trying to create a hybrid MD simulation with aimnet2 for pdbid: 3WJJ where part of the system is simulated using the NNP.

`
import openmm as mm
from openmm import *
from openmm.app import *
from openmm.unit import *
from parmed.openmm.reporters import *
from tqdm import tqdm
import pytraj as pt
from sys import stdout, exit, stderr, argv
import os, math, fnmatch
import argparse
import time
from time import sleep
import warnings
import numpy as np
from openmmml import MLPotential

warnings.filterwarnings('ignore')

parser = argparse.ArgumentParser()
parser.add_argument("-p","--topology",help="Amber topology file",required=True)
parser.add_argument("-c","--coordinates",help="Amber coordinates file",required=True)
parser.add_argument("-s","--structure",help="PDB file",required=True)
parser.add_argument("-temp","--temperature",help="Temperature (default: 300K)", type=int, default=300, required=False)
parser.add_argument("-nvt","--nvt",help="Number of steps NVT heating",required=True)
parser.add_argument("-npt","--npt",help="Number of steps NPT equilibration",required=True)
parser.add_argument("-prod","--production",help="Number of steps NPT production",required=True)
parser.add_argument("-st","--strides",help="Number of strides for NPT production",required=True)
parser.add_argument("-mlmodel","--mlmodel",help="ML model name (default: ani2x)",default="ani2x",required=False)
parser.add_argument("-gpuID",help="GPU ID(s) (default: '0,1' for both GPUs)", default="1", required=False)

args = parser.parse_args()

time_nvt_1fs = (float(args.nvt) * 1000 *0.2) / 0.001
time_nvt_2fs = (float(args.nvt) * 1000 *0.8) / 0.002
time_npt = (float(args.npt) * 1000) / 0.002
time_production = (float(args.production) * 1000) / 0.002

def apply_restraints(system, pdb, stage="heating"):
"""
Apply position restraints following md_simulation.py approach.

Args:
    system: OpenMM System object
    pdb: PDB file with structure
    stage: Simulation stage ("heating" or "equilibration")
    
Returns:
    Modified system with restraints
"""
print("\t- Applying position restraints on protein CA atoms and non-protein heavy atoms...")

# Create restraint force
restraint = CustomExternalForce("k*periodicdistance(x, y, z, x0, y0, z0)^2")

# Set initial restraint strength based on stage
initial_k = 10.0 if stage == "heating" else 0.0  # 10 for heating, 0 for equilibration
restraint.addGlobalParameter("k", initial_k * kilocalories_per_mole / angstroms**2)

restraint.addPerParticleParameter("x0")
restraint.addPerParticleParameter("y0")
restraint.addPerParticleParameter("z0")

restrained_atoms = 0
# Add restraints matching md_simulation.py logic
for res in pdb.topology.residues():
    if res.name not in ["WAT", "HOH", "SOL", "Na+", "Cl-"]:
        # Check if residue has CA or N atoms (protein residue)
        if any([x.name for x in res.atoms() if x.name in ["CA", "N"]]):
            for atom in res.atoms():
                if atom.name == "CA":  # Only restrain CA atoms for protein
                    restraint.addParticle(atom.index, pdb.positions[atom.index])
                    restrained_atoms += 1
        else:  # Non-protein heavy atoms (e.g., ligand)
            for atom in res.atoms():
                if "H" not in atom.name:  # Skip hydrogens
                    restraint.addParticle(atom.index, pdb.positions[atom.index])
                    restrained_atoms += 1

system.addForce(restraint)
print(f"\t- Added restraints to {restrained_atoms} atoms")
return system

def create_mixed_system(prmtop, pdb, is_minimization=False):
"""
Create a mixed MM/ML system using OpenMM-ML with AIMNet2 and EMLE for electrostatic embedding.

Args:
    prmtop: AmberPrmtopFile for the system
    pdb: PDB file with structure
    is_minimization: Whether this system is for minimization (default: False)

Returns:
    A System object with mixed MM/ML potentials
"""
print("\t- Creating mixed MM/ML system with AIMNet2 and EMLE embedding")

# Define the ML region based on specific residues in chains A
ml_residues = {
    'A': list(range(233, 249))
}
# Check for the presence of chains A in the PDB
chains_in_pdb = {chain.id: chain for chain in pdb.topology.chains()}

# Initialize the list of atoms in the ML region
ml_atoms = []
target_renum_ids = []
# Identify atoms directly from the PDB if chain A is present
for chain_id, residues in ml_residues.items():
    if chain_id in chains_in_pdb:  # Chain is present in the PDB
        for atom in pdb.topology.atoms():
            if atom.residue.chain.id == chain_id and atom.residue.id in residues:
                ml_atoms.append(atom.index)
        print(f"\t- Found atoms in chain {chain_id} directly from the PDB.")
    else:
        print(f"\t- Chain {chain_id} not found in the PDB. Checking renumbered residues...")

# If no ML atoms are found, open the cleaned_renum.txt to find renumbered residues
if not ml_atoms:
    renum_file = "cleaned_renum.txt"  # Path to the renumbered residues file
    if os.path.exists(renum_file):
        # First pass: Read the renum file and map original residues to renumbered residues
        try:
            with open(renum_file, 'r') as f:
                for line in f:
                    parts = line.strip().split()
                    if len(parts) >= 5:
                        chain_id = parts[1]
                        orig_id = int(parts[2])
                        renum_id = int(parts[4])

                        # Check if this residue is in our ml_residues
                        for chain, residues in ml_residues.items():
                            if chain == chain_id and orig_id in residues:
                                target_renum_ids.append(renum_id)
                                print(f"\t  Found renumbered ID {renum_id} for residue {chain}:{orig_id}")

            for atom in pdb.topology.atoms():
                try:
                    res_id = int(atom.residue.id)
                    if res_id in target_renum_ids:
                        ml_atoms.append(atom.index)
                except ValueError:
                    continue

            print(f"\t- Found {len(ml_atoms)} atoms in ML region using renumbering file")

        except Exception as e:
            print(f"\t- Error processing renumbering file: {str(e)}")

# If still no ML atoms found, use MM-only system
if not ml_atoms:
    print("\t- Warning: No atoms found for the specified ML region")
    print("\t- Using MM-only system")
    system_kwargs = {
        'nonbondedMethod': PME,
        'nonbondedCutoff': 1.0 * nanometers,
        'constraints': HBonds
    }
    if is_minimization:
        system_kwargs['ewaldErrorTolerance'] = 0.001
    else:
        system_kwargs['rigidWater'] = False

    return prmtop.createSystem(**system_kwargs)

print(f"\t- Found {len(ml_atoms)} atoms in the ML region")

# Create base MM system from prmtop
system_kwargs = {
    'nonbondedMethod': PME,
    'nonbondedCutoff': 1.0 * nanometers,
    'constraints': HBonds,
    'removeCMMotion': True,
    'hydrogenMass': 3.024 * amu
}

if is_minimization:
    system_kwargs['ewaldErrorTolerance'] = 0.001
else:
    system_kwargs['rigidWater'] = False

# Create the MM system
mm_system = prmtop.createSystem(**system_kwargs)

# Initialize the ML potential (AIMNet2)
model_path = '/root/miniconda3/envs/workflow2/lib/python3.10/site-packages/openmmml/models/aimnet2_wb97m-d3_ens.jpt'
potential = MLPotential('aimnet2-wb97m-d3', model=model_path)

# Create the mixed ML/MM system with AIMNet2
# Add charge=0 to satisfy the AIMNet2 implementation requirement
hybrid_system = potential.createMixedSystem(
    pdb.topology, mm_system, ml_atoms, charge=0
)

print(f"\t- Successfully created hybrid AIMNet2/MM system with optimized EMLE settings")
return hybrid_system

`