Bug fix

tingiskhan · Jul 14, 2024 · 831cebe · 831cebe
1 parent 0747d38
commit 831cebe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/numpyro_sts/smooth_llt.py b/numpyro_sts/smooth_llt.py
@@ -24,6 +24,6 @@ def __init__(self, n: int, std: ArrayLike, initial_value: ArrayLike, **kwargs):
         matrix = jnp.array([[1.0, 1.0], [0.0, 1.0]])
         offset = jnp.zeros_like(initial_value)
 
-        column_mask = np.eye(2, 1, dtype=np.bool_).squeeze(-1)
+        column_mask = np.array([False, True])
 
         super().__init__(n, offset, matrix, std, initial_value, column_mask=column_mask, **kwargs)