From f7efd394431b2e461079d0cec987d2d735d8cf67 Mon Sep 17 00:00:00 2001
From: Yerdos Ordabayev <yordabay@broadinstitute.org>
Date: Fri, 7 Jul 2023 15:33:05 +0000
Subject: [PATCH] fix weakref

---
 tapqir/models/hmm.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tapqir/models/hmm.py b/tapqir/models/hmm.py
index e71933d..0dfb36c 100644
--- a/tapqir/models/hmm.py
+++ b/tapqir/models/hmm.py
@@ -670,4 +670,6 @@ def z_sample(self, num_samples):
     def compute_params(self, CI):
         params = super().compute_params(CI)
         params["z_trans"] = pyro.param("z_trans").cpu()
+        # Remove weakrefs in preparation for pickling.
+        params["z_trans"].__dict__.pop("unconstrained", None)
         return params