Update app.py
Browse files
app.py
CHANGED
@@ -13,6 +13,7 @@ def load_model(repo_name):
|
|
13 |
tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
|
14 |
generator_conf = GenerationConfig.from_pretrained(repo_name)
|
15 |
model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True, torch_dtype=torch.bfloat16)
|
|
|
16 |
return tokenizer, generator_conf, model
|
17 |
|
18 |
tokenizer, generator_conf, model = load_model(REPO_NAME)
|
|
|
13 |
tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
|
14 |
generator_conf = GenerationConfig.from_pretrained(repo_name)
|
15 |
model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True, torch_dtype=torch.bfloat16)
|
16 |
+
model.to('cuda')
|
17 |
return tokenizer, generator_conf, model
|
18 |
|
19 |
tokenizer, generator_conf, model = load_model(REPO_NAME)
|