logo
Browse Source

Fix model name issue

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 1 year ago
parent
commit
93ed7f5e5e
  1. 5
      auto_transformers.py

5
auto_transformers.py

@ -81,12 +81,11 @@ class AutoTransformers(NNOperator):
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
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
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(
model_name=self.model_name, device=self.device, checkpoint_path=checkpoint_path)
if tokenizer is None:

Loading…
Cancel
Save