fix weakref

gelles-brandeis · Jul 7, 2023 · f7efd39 · f7efd39
1 parent 0292aad
commit f7efd39
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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