bryanmildort commited on
Commit
a9f3afd
·
1 Parent(s): 7e21c92

Update app.py

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