diff --git a/Dockerfile b/Dockerfile
index 0c32dd7..062bfe6 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,3 +2,4 @@ FROM ssages/pysages-hoomd:latest
 
 COPY . openmm-dlext
 RUN cd openmm-dlext && mkdir build && cd build && cmake .. && make install
+RUN python -c "import openmm_dlext"
diff --git a/wrappers/python/DLExtForce.i b/wrappers/python/DLExtForce.i
index cd94089..359daea 100644
--- a/wrappers/python/DLExtForce.i
+++ b/wrappers/python/DLExtForce.i
@@ -1,6 +1,15 @@
 %module dlpack_extension_swig
 
-%import(module = "simtk.openmm") "OpenMMForce.i"
+%pythoncode
+%{
+from importlib import import_module
+try:
+    openmm = import_module("openmm")
+except:
+    openmm = import_module("simtk.openmm")
+%}
+
+%import(module = "openmm") "OpenMMForce.i"
 %include "std_string.i"
 
 %{