From 2dc3b8d7c3b812146de154642a61b5b2fe49a61a Mon Sep 17 00:00:00 2001
From: Pablo Zubieta <8410335+pabloferz@users.noreply.github.com>
Date: Thu, 3 Aug 2023 16:12:02 -0500
Subject: [PATCH] Make sure that lammps contexts are properly finalized

---
 pysages/backends/lammps.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pysages/backends/lammps.py b/pysages/backends/lammps.py
index aa07a52a..376254e7 100644
--- 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