fix erase

InternLM · Dec 18, 2023 · 9f491a8 · 9f491a8
1 parent 20d8f47
commit 9f491a8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/turbomind/models/llama/SequenceManager.cc b/src/turbomind/models/llama/SequenceManager.cc
@@ -61,7 +61,7 @@ bool SequenceManager::Contains(uint64_t id)
     return sequences_.find(id) != sequences_.end();
 }
 
-void SequenceManager::Erase(std::map<uint64_t, Sequence>::iterator it)
+void SequenceManager::Erase(std::map<uint64_t, Sequence>::iterator& it)
 {
     auto& seq = it->second;
     if (seq.status == Sequence::kCached) {
@@ -72,7 +72,7 @@ void SequenceManager::Erase(std::map<uint64_t, Sequence>::iterator it)
         UpdateAndSetUnlock(seq);
     }
     freed_.insert(freed_.end(), seq.blocks.begin(), seq.blocks.end());
-    sequences_.erase(it);
+    it = sequences_.erase(it);
 }
 
 bool SequenceManager::Erase(uint64_t id)

diff --git a/src/turbomind/models/llama/SequenceManager.h b/src/turbomind/models/llama/SequenceManager.h
@@ -107,7 +107,7 @@ class SequenceManager {
     }
 
 private:
-    void Erase(std::map<uint64_t, Sequence>::iterator it);
+    void Erase(std::map<uint64_t, Sequence>::iterator& it);
 
     void CommitUnlockAndFree();