diff --git a/examples/pose_graph_simple.py b/examples/pose_graph_simple.py
index 436c26c..b0829c9 100644
--- a/examples/pose_graph_simple.py
+++ b/examples/pose_graph_simple.py
@@ -48,7 +48,7 @@
 graph = jaxls.FactorGraph.make(factors, vars)
 
 # Solve the optimization problem.
-solution = graph.solve(linear_solver=jaxls.DenseLinearSolver())
+solution = graph.solve()
 print("All solutions", solution)
 print("Pose 0", solution[vars[0]])
 print("Pose 1", solution[vars[1]])
diff --git a/src/jaxls/_preconditioning.py b/src/jaxls/_preconditioning.py
index 9f9c9fe..2368740 100644
--- a/src/jaxls/_preconditioning.py
+++ b/src/jaxls/_preconditioning.py
@@ -25,9 +25,9 @@ def make_point_jacobi_precoditioner(
         block_l2_cols = jnp.sum(block_row.blocks_concat**2, axis=1).flatten()
         indices = jnp.concatenate(
             [
-                (start_col[:, None] + jnp.arange(width)[None, :])
-                for start_col, width in zip(
-                    block_row.start_cols, block_row.block_widths
+                (start_col[:, None] + jnp.arange(block_cols)[None, :])
+                for start_col, block_cols in zip(
+                    block_row.start_cols, block_row.block_num_cols
                 )
             ],
             axis=1,