Browse Source
Fix model name issue
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
2 additions and
3 deletions
-
auto_transformers.py
|
|
@ -81,12 +81,11 @@ class AutoTransformers(NNOperator): |
|
|
if device is None: |
|
|
if device is None: |
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
self.device = device |
|
|
self.device = device |
|
|
|
|
|
|
|
|
model_list = self.supported_model_names() |
|
|
|
|
|
assert model_name in model_list, f"Invalid model name: {model_name}. Supported model names: {model_list}" |
|
|
|
|
|
self.model_name = model_name |
|
|
self.model_name = model_name |
|
|
|
|
|
|
|
|
if self.model_name: |
|
|
if self.model_name: |
|
|
|
|
|
model_list = self.supported_model_names() |
|
|
|
|
|
assert model_name in model_list, f"Invalid model name: {model_name}. Supported model names: {model_list}" |
|
|
self.model = Model( |
|
|
self.model = Model( |
|
|
model_name=self.model_name, device=self.device, checkpoint_path=checkpoint_path) |
|
|
model_name=self.model_name, device=self.device, checkpoint_path=checkpoint_path) |
|
|
if tokenizer is None: |
|
|
if tokenizer is None: |
|
|
|
