ejschwartz commited on
Commit
3cb2b56
·
1 Parent(s): 33e5ee4

try moving to cuda again

Browse files
Files changed (1) hide show
  1. app.py +2 -4
app.py CHANGED
@@ -18,13 +18,11 @@ tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoderbase-3b")
18
  vardecoder_model = AutoModelForCausalLM.from_pretrained(
19
  "ejschwartz/resym-vardecoder",
20
  torch_dtype=torch.bfloat16,
21
- device_map="auto"
22
- )
23
  fielddecoder_model = AutoModelForCausalLM.from_pretrained(
24
  "ejschwartz/resym-fielddecoder",
25
  torch_dtype=torch.bfloat16,
26
- device_map="auto"
27
- )
28
 
29
  make_gradio_client = lambda: Client("https://ejschwartz-resym-field-helper.hf.space/")
30
 
 
18
  vardecoder_model = AutoModelForCausalLM.from_pretrained(
19
  "ejschwartz/resym-vardecoder",
20
  torch_dtype=torch.bfloat16,
21
+ ).to("cuda")
 
22
  fielddecoder_model = AutoModelForCausalLM.from_pretrained(
23
  "ejschwartz/resym-fielddecoder",
24
  torch_dtype=torch.bfloat16,
25
+ ).to("cuda")
 
26
 
27
  make_gradio_client = lambda: Client("https://ejschwartz-resym-field-helper.hf.space/")
28