1 files changed, 6 insertions, 2 deletions
diff --git a/search/edge.hh b/search/edge.hh
index 2cdc05f37..3b096c1ae 100644
--- a/search/edge.hh
+++ b/search/edge.hh
@@ -62,12 +62,16 @@ class PartialEdge {
 
 class PartialEdgePool {
   public:
-    PartialEdge Allocate(Arity arity) {
+    PartialEdge Allocate(Arity arity, Arity chart_states) {
       return PartialEdge(
-          pool_.Allocate(sizeof(Score) + sizeof(Arity) + arity * sizeof(PartialVertex) + (arity + 1) * sizeof(lm::ngram::ChartState)),
+          pool_.Allocate(sizeof(Score) + sizeof(Arity) + arity * sizeof(PartialVertex) + chart_states * sizeof(lm::ngram::ChartState)),
           arity);
     }
 
+    PartialEdge Allocate(Arity arity) {
+      return Allocate(arity, arity + 1);
+    }
+
   private:
     util::Pool pool_;
 };