Make sure that lammps contexts are properly finalized

SSAGESLabs · Aug 3, 2023 · 2dc3b8d · 2dc3b8d
1 parent 70640aa
commit 2dc3b8d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/pysages/backends/lammps.py b/pysages/backends/lammps.py
@@ -7,6 +7,7 @@
 """
 
 import importlib
+import weakref
 from functools import partial
 
 import jax
@@ -247,5 +248,7 @@ def bind(sampling_context: SamplingContext, callback: Optional[Callable], **kwar
  # Unfortunately, the default implementation of `lammps.__exit__` closes
  # the lammps instance, so we need to overwrite it.
  context.__exit__ = lambda *args: None
+ # Ensure that the lammps context is properly finalized.
+ weakref.finalize(context, context.finalize)
 
  return sampler