Torch2 compatibility (#106)

* Update ci

* port the latest CI versions to the self-hosted runner

* forgot to remove bit from version I copy

* see if cuda 11.7 works

* Update swig typecheck to use py::isinstance instead of as_module

* run CUDA tests on GPU runner

* Update TorchForce swig input typemap to avoid as_module

* Fix formatting

* Switch to using micromamba

* Make sure temporary files are deleted even if there was an error

* Update swig typemaps to not use a temp file

* Revert "Merge remote-tracking branch 'origin/feat/add_aws_gpu_runner' into torch2"

This reverts commit 0a58de4, reversing
changes made to a1c16c6.

* run gpu tests on this branch

* actually test with pytorch2

* Revert "actually test with pytorch2"

This reverts commit 4e7a4fb.

* Revert "Merge remote-tracking branch 'origin/feat/add_aws_gpu_runner' into torch2"

This reverts commit 0a58de4, reversing
changes made to a1c16c6.

---------

Co-authored-by: Mike Henry <11765982+mikemhenry@users.noreply.github.com>

.github/workflows/CI.yml

@@ -33,13 +33,13 @@ jobs:
  pytorch-version: "1.11.*"
 
  # Latest supported versions
- - name: Linux (CUDA 11.2, Python 3.10, PyTorch 1.12)
+ - name: Linux (CUDA 11.8, Python 3.10, PyTorch 2.0)
  os: ubuntu-22.04
- cuda-version: "11.2.2"
+ cuda-version: "11.8.0"
  gcc-version: "10.3.*"
- nvcc-version: "11.2"
+ nvcc-version: "11.8"
  python-version: "3.10"
- pytorch-version: "1.12.*"
+ pytorch-version: "2.0.*"
 
  - name: MacOS (Python 3.9, PyTorch 1.9)
  os: macos-11

python/openmmtorch.i

@@ -14,6 +14,7 @@
 #include "openmm/RPMDIntegrator.h"
 #include "openmm/RPMDMonteCarloBarostat.h"
 #include <torch/csrc/jit/python/module_python.h>
+#include <torch/csrc/jit/serialization/import.h>
 %}
 
 /*
@@ -28,23 +29,27 @@
  }
 }
 
-%typemap(in) const torch::jit::Module&(torch::jit::Module module) {
+%typemap(in) const torch::jit::Module&(torch::jit::Module mod) {
  py::object o = py::reinterpret_borrow<py::object>($input);
- module = torch::jit::as_module(o).value();
- $1 = &module;
+ py::object pybuffer = py::module::import("io").attr("BytesIO")();
+ py::module::import("torch.jit").attr("save")(o, pybuffer);
+ std::string s = py::cast<std::string>(pybuffer.attr("getvalue")());
+ std::stringstream buffer(s);
+ mod = torch::jit::load(buffer);
+ $1 = &mod;
 }
 
 %typemap(out) const torch::jit::Module& {
- auto fileName = std::tmpnam(nullptr);
- $1->save(fileName);
- $result = py::module::import("torch.jit").attr("load")(fileName).release().ptr();
- //This typemap assumes that torch does not require the file to exist after construction
- std::remove(fileName);
+ std::stringstream buffer;
+ $1->save(buffer);
+ auto pybuffer = py::module::import("io").attr("BytesIO")(py::bytes(buffer.str()));
+ $result = py::module::import("torch.jit").attr("load")(pybuffer).release().ptr();
 }
 
 %typecheck(SWIG_TYPECHECK_POINTER) const torch::jit::Module& {
  py::object o = py::reinterpret_borrow<py::object>($input);
- $1 = torch::jit::as_module(o).has_value() ? 1 : 0;
+ py::handle ScriptModule = py::module::import("torch.jit").attr("ScriptModule");
+ $1 = py::isinstance(o, ScriptModule);
 }
 
 namespace std {