Unverified Commit 0731668c authored by Gao, Xiang's avatar Gao, Xiang Committed by GitHub
Browse files

nvidia profiling for md (#411)

* nvidia profiling for md

* 1 model
parent 2de2ecbc
import ase
import ase.io
import ase.md
import argparse
import torchani
import autonvtx
import torch
parser = argparse.ArgumentParser()
parser.add_argument('filename', help="file for the molecule")
args = parser.parse_args()
molecule = ase.io.read(args.filename)
model = torchani.models.ANI1x()[0].cuda()
calculator = model.ase()
molecule.set_calculator(calculator)
dyn = ase.md.verlet.VelocityVerlet(molecule, timestep=1 * ase.units.fs)
dyn.run(1000) # warm up
torch.cuda.cudart().cudaProfilerStart()
autonvtx(model)
with torch.autograd.profiler.emit_nvtx(record_shapes=True):
dyn.run(10)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment