ejschwartz commited on
Commit
d3fb139
·
1 Parent(s): d5e44b4

More experimenting

Browse files
Files changed (1) hide show
  1. app.py +6 -2
app.py CHANGED
@@ -22,11 +22,15 @@ tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoderbase-3b")
22
  vardecoder_model = AutoModelForCausalLM.from_pretrained(
23
  "ejschwartz/resym-vardecoder",
24
  torch_dtype=torch.bfloat16,
25
- ).to("cuda")
 
 
26
  fielddecoder_model = AutoModelForCausalLM.from_pretrained(
27
  "ejschwartz/resym-fielddecoder",
28
  torch_dtype=torch.bfloat16,
29
- ).to("cuda")
 
 
30
 
31
  make_gradio_client = lambda: Client("https://ejschwartz-resym-field-helper.hf.space/")
32
 
 
22
  vardecoder_model = AutoModelForCausalLM.from_pretrained(
23
  "ejschwartz/resym-vardecoder",
24
  torch_dtype=torch.bfloat16,
25
+ low_cpu_mem_usage=True,
26
+ device_map="cuda",
27
+ )
28
  fielddecoder_model = AutoModelForCausalLM.from_pretrained(
29
  "ejschwartz/resym-fielddecoder",
30
  torch_dtype=torch.bfloat16,
31
+ low_cpu_mem_usage=True,
32
+ device_map="cuda",
33
+ )
34
 
35
  make_gradio_client = lambda: Client("https://ejschwartz-resym-field-helper.hf.space/")
36