Fix crash

vicariousinc · Oct 20, 2021 · 84672d4 · 84672d4
1 parent 7925ea7
commit 84672d4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/pgmax/fg/graph.py b/pgmax/fg/graph.py
@@ -614,13 +614,12 @@ def __setitem__(
                     .set(evidence_val)
                 )
         else:
-            start_index = self.factor_graph._vars_to_starts[
-                self.factor_graph._variable_group[key]
-            ]
+            var = self.factor_graph._variable_group[key]
+            start_index = self.factor_graph._vars_to_starts[var]
             self.value = (
                 jax.device_put(self.value)
                 .at[start_index : start_index + var.num_states]
-                .set(evidence_val)
+                .set(evidence)
             )
 
 

diff --git a/tests/fg/test_graph.py b/tests/fg/test_graph.py
@@ -19,7 +19,7 @@ def test_onevar_graph():
 
     assert "Unsupported default message mode" in str(verror.value)
     with pytest.raises(ValueError) as verror:
-        graph.Evidence(factor_graph=fg, default_mode="zeros", init_value=np.zeros(1))
+        graph.Evidence(factor_graph=fg, default_mode="zeros", value=np.zeros(1))
 
     assert "Should specify only" in str(verror.value)
     with pytest.raises(ValueError) as verror: