diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl
index 3150400d..81ef80ba 100644
--- a/src/caches/beliefpropagationcache.jl
+++ b/src/caches/beliefpropagationcache.jl
@@ -13,7 +13,7 @@ end
 default_partitioning(ψ::AbstractITensorNetwork) = group(v -> v, vertices(ψ))
 
 #We could probably do something cleverer here based on graph partitioning algorithms: https://en.wikipedia.org/wiki/Graph_partition.
-default_partitioning(f::AbstractFormNetwork) = group(v -> first(v), vertices(f))
+default_partitioning(f::AbstractFormNetwork) = group(v -> state_vertex(f, v), vertices(f))
 default_cache_update_kwargs(cache) = (; maxiter=20, tol=1e-5)
 
 function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor})
diff --git a/src/formnetworks/abstractformnetwork.jl b/src/formnetworks/abstractformnetwork.jl
index a8187751..82d79a24 100644
--- a/src/formnetworks/abstractformnetwork.jl
+++ b/src/formnetworks/abstractformnetwork.jl
@@ -75,3 +75,4 @@ inv_vertex_map(f::AbstractFormNetwork) = v -> first(v)
 operator_vertex(f::AbstractFormNetwork, v) = operator_vertex_map(f)(v)
 bra_vertex(f::AbstractFormNetwork, v) = bra_vertex_map(f)(v)
 ket_vertex(f::AbstractFormNetwork, v) = ket_vertex_map(f)(v)
+state_vertex(f::AbstractFormNetwork, v) = inv_vertex_map(f)(v)