Browse Source
Optimize triton
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
2 additions and
2 deletions
-
auto_transformers.py
|
|
@ -128,12 +128,12 @@ class AutoTransformers(NNOperator): |
|
|
log.error(f'Fail to tokenize inputs: {e}') |
|
|
log.error(f'Fail to tokenize inputs: {e}') |
|
|
raise e |
|
|
raise e |
|
|
try: |
|
|
try: |
|
|
outs = self.model(**inputs) |
|
|
|
|
|
|
|
|
outs = self.model(**inputs).to('cpu') |
|
|
except Exception as e: |
|
|
except Exception as e: |
|
|
log.error(f'Invalid input for the model: {self.model_name}') |
|
|
log.error(f'Invalid input for the model: {self.model_name}') |
|
|
raise e |
|
|
raise e |
|
|
outs = self.post_proc(outs, inputs) |
|
|
outs = self.post_proc(outs, inputs) |
|
|
features = outs.cpu().detach().numpy() |
|
|
|
|
|
|
|
|
features = outs.detach().numpy() |
|
|
if isinstance(data, str): |
|
|
if isinstance(data, str): |
|
|
features = features.squeeze(0) |
|
|
features = features.squeeze(0) |
|
|
else: |
|
|
else: |
|
|
|
