Fix memory leak when using scan (#1469)

pyro-ppl · Aug 14, 2022 · 0e63cfa · 0e63cfa
1 parent 589b352
commit 0e63cfa
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/numpyro/ops/pytree.py b/numpyro/ops/pytree.py
@@ -32,6 +32,12 @@ def tree_flatten(self):
                                 # set to None to avoid leaks during tracing by JAX
                                 kwargs["rng_key"] = None
                             aux_trace[name][key] = kwargs
+                        elif key == "infer":
+                            kwargs = site["infer"].copy()
+                            if "_scan_current_index" in kwargs:
+                                # set to None to avoid leaks during tracing by JAX
+                                kwargs["_scan_current_index"] = None
+                            aux_trace[name][key] = kwargs
                         else:
                             aux_trace[name][key] = site[key]
         # keep the site order information because in JAX, flatten and unflatten do not preserve