[LLM] fix bloom and starcoder memory release (#8728)

This commit is contained in:
Yishuo Wang 2023-08-11 11:18:19 +08:00 committed by GitHub
parent 33d9ad234f
commit 3d5a7484a2
2 changed files with 10 additions and 6 deletions

View file

@ -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

View file

@ -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