Commit
·
a9f3afd
1
Parent(s):
7e21c92
Update app.py
Browse files
app.py
CHANGED
@@ -17,7 +17,7 @@ import torch
|
|
17 |
# device_map = infer_auto_device_map(model, dtype="float16")
|
18 |
# st.write(device_map)
|
19 |
|
20 |
-
model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes_summary", low_cpu_mem_usage=True)
|
21 |
tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes_summary")
|
22 |
# model = model.to(device)
|
23 |
|
|
|
17 |
# device_map = infer_auto_device_map(model, dtype="float16")
|
18 |
# st.write(device_map)
|
19 |
|
20 |
+
model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes_summary", low_cpu_mem_usage=True).cuda()
|
21 |
tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes_summary")
|
22 |
# model = model.to(device)
|
23 |
|