[LLM] fix bloom and starcoder memory release (#8728)
This commit is contained in:
parent
33d9ad234f
commit
3d5a7484a2
2 changed files with 10 additions and 6 deletions
|
|
@ -313,9 +313,6 @@ class Bloom(GenerationMixin):
|
|||
}
|
||||
}
|
||||
|
||||
def free(self):
|
||||
bloom_free(self.ctx)
|
||||
|
||||
def _tokenize(self, text: bytes, add_bos: bool = False) -> List[int]:
|
||||
"""Tokenize a string.
|
||||
|
||||
|
|
@ -427,3 +424,8 @@ class Bloom(GenerationMixin):
|
|||
seed=self.seed,
|
||||
n_threads=self.n_threads,
|
||||
n_batch=self.n_batch)
|
||||
|
||||
def __del__(self):
|
||||
if self.ctx is not None:
|
||||
bloom_free(self.ctx)
|
||||
self.ctx = None
|
||||
|
|
|
|||
|
|
@ -315,9 +315,6 @@ class Starcoder(GenerationMixin):
|
|||
}
|
||||
}
|
||||
|
||||
def free(self):
|
||||
starcoder_free(self.ctx)
|
||||
|
||||
def _tokenize(self, text: bytes, add_bos: bool = False) -> List[int]:
|
||||
"""Tokenize a string.
|
||||
|
||||
|
|
@ -433,3 +430,8 @@ class Starcoder(GenerationMixin):
|
|||
seed=self.seed,
|
||||
n_threads=self.n_threads,
|
||||
n_batch=self.n_batch)
|
||||
|
||||
def __del__(self):
|
||||
if self.ctx is not None:
|
||||
starcoder_free(self.ctx)
|
||||
self.ctx = None
|
||||
|
|
|
|||
Loading…
Reference in a new issue