schuler commited on
Commit
4f41458
·
verified ·
1 Parent(s): 7b5b897

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -0
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)