Hugging Face's logo

Spaces:
joaogante
/
assisted_generation_demo
Running on Zero

App Files Community
11
joaogante HF staff commited on
Commit
a304cf8
·
verified ·
1 Parent(s): 4ffe071

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -10,7 +10,7 @@ model_id = "facebook/opt-6.7b"
10
  assistant_id = "facebook/opt-125m"
11
 
12
  model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
13
- assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id, torch_dtype=torch.bfloat16, device_map="auto")
14
  tokenizer = AutoTokenizer.from_pretrained(model_id)
15
 
16
  @spaces.GPU
 
10
  assistant_id = "facebook/opt-125m"
11
 
12
  model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
13
+ assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.bfloat16)
14
  tokenizer = AutoTokenizer.from_pretrained(model_id)
15
 
16
  @spaces.GPU