HammXG commited on
Commit
19939f6
Β·
verified Β·
1 Parent(s): c17cea9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -9,9 +9,10 @@ tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
9
  model = AutoModelForCausalLM.from_pretrained(
10
  model_id,
11
  device_map="auto",
12
- torch_dtype=torch.float16,
13
  trust_remote_code=True,
14
- revision="main"
 
 
15
  )
16
 
17
  pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
9
  model = AutoModelForCausalLM.from_pretrained(
10
  model_id,
11
  device_map="auto",
 
12
  trust_remote_code=True,
13
+ revision="main",
14
+ torch_dtype=torch.float16,
15
+ low_cpu_mem_usage=True
16
  )
17
 
18
  pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)