Unable to run dwell time analysis with the new posterior sample option

Running 1.1.15, default cosmos+hmm settings
here is the error:
Dwell-time analysis ...
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/anaconda3/envs/tapqir-env/lib/python3.8/site-packages/tapqir/gui.py:1350, in dtCmd(b, layout, out)
   1348 with out:
   1349     logger.info("Dwell-time analysis ...")
-> 1350     dwelltime(
   1351         **layout.kwargs,
   1352         progress_bar=tqdm_notebook,
   1353     )
   1354     logger.info("Dwell-time analysis: Done")
   1355 out.clear_output(wait=True)

File ~/anaconda3/envs/tapqir-env/lib/python3.8/site-packages/tapqir/main.py:1213, in dwelltime(model, K, cuda, num_samples, num_iter, progress_bar)
   1211     logger.exception(f"Failed to load {err.name} file")
   1212     return 1
-> 1213 z_samples = model.z_sample(num_samples=num_samples)
   1214 z_samples_masked = z_samples[:, model.data.mask[: model.data.N]]
   1215 for c in range(model.data.C):

File ~/anaconda3/envs/tapqir-env/lib/python3.8/site-packages/tapqir/models/hmm.py:659, in hmm.z_sample(self, num_samples)
    658 def z_sample(self, num_samples):
--> 659     init_probs = pyro.param("z_trans").data[: self.data.N, 0, :, 0]
    660     init_probs = init_probs.expand((num_samples,) + init_probs.shape)
    661     x = dist.Categorical(init_probs).sample()

File ~/anaconda3/envs/tapqir-env/lib/python3.8/site-packages/pyro/primitives.py:78, in param(name, init_tensor, constraint, event_dim)
     76 # Note effectful(-) requires the double passing of name below.
     77 args = (name,) if init_tensor is None else (name, init_tensor)
---> 78 return _param(*args, constraint=constraint, event_dim=event_dim, name=name)

File ~/anaconda3/envs/tapqir-env/lib/python3.8/site-packages/pyro/poutine/runtime.py:264, in effectful.<locals>._fn(*args, **kwargs)
    261 is_observed = value is not None
    263 if not am_i_wrapped():
--> 264     return fn(*args, **kwargs)
    265 else:
    266     msg = {
    267         "type": type,
    268         "name": name,
   (...)
    280         "infer": infer,
    281     }

File ~/anaconda3/envs/tapqir-env/lib/python3.8/site-packages/pyro/params/param_store.py:215, in ParamStoreDict.get_param(self, name, init_tensor, constraint, event_dim)
    199 """
    200 Get parameter from its name. If it does not yet exist in the
    201 ParamStore, it will be created and stored.
   (...)
    212 :rtype: torch.Tensor
    213 """
    214 if init_tensor is None:
--> 215     return self[name]
    216 else:
    217     return self.setdefault(name, init_tensor, constraint)

File ~/anaconda3/envs/tapqir-env/lib/python3.8/site-packages/pyro/params/param_store.py:106, in ParamStoreDict.__getitem__(self, name)
    102 def __getitem__(self, name):
    103     """
    104     Get the *constrained* value of a named parameter.
    105     """
--> 106     unconstrained_value = self._params[name]
    108     # compute the constrained value
    109     constraint = self._constraints[name]

KeyError: 'z_trans
gelles-brandeis / tapqir

Unable to run dwell time analysis with the new posterior sample option #408
