torchmd-net icon indicating copy to clipboard operation
torchmd-net copied to clipboard

Published 20 hours ago • verified publisher icon torchmd

POC: optimize the graph network with CUDA Graph

Open raimis opened this issue 3 years ago • 0 comments

This is just a proof of concept. DO NOT MERGE!

  • [x] Implement DistanceBruteForce
  • [x] Integrate into TorchMD_GN
  • [x] Make compatible with CUDA Graphs
  • [x] Tests
  • [x] Benchmark

~~Depends on #59~~

raimis avatar Feb 23 '22 14:02 raimis