ejschwartz commited on
Commit
762a224
·
1 Parent(s): 84b4a86
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -16,10 +16,10 @@ huggingface_hub.login(token=hf_key)
16
  tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoderbase-3b")
17
  vardecoder_model = AutoModelForCausalLM.from_pretrained(
18
  "ejschwartz/resym-vardecoder", torch_dtype=torch.bfloat16#, device_map={"": 0}
19
- )
20
  fielddecoder_model = AutoModelForCausalLM.from_pretrained(
21
  "ejschwartz/resym-fielddecoder", torch_dtype=torch.bfloat16#, device_map={"": 0}
22
- )
23
 
24
  gradio_client = Client("https://ejschwartz-resym-field-helper.hf.space/")
25
 
 
16
  tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoderbase-3b")
17
  vardecoder_model = AutoModelForCausalLM.from_pretrained(
18
  "ejschwartz/resym-vardecoder", torch_dtype=torch.bfloat16#, device_map={"": 0}
19
+ ).to("cuda")
20
  fielddecoder_model = AutoModelForCausalLM.from_pretrained(
21
  "ejschwartz/resym-fielddecoder", torch_dtype=torch.bfloat16#, device_map={"": 0}
22
+ ).to("cuda")
23
 
24
  gradio_client = Client("https://ejschwartz-resym-field-helper.hf.space/")
25