zhaozitian commited on
Commit
ba83313
·
1 Parent(s): c6bce9b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -27,7 +27,7 @@ except:
27
  if device == "cuda":
28
  model = LlamaForCausalLM.from_pretrained(
29
  BASE_MODEL,
30
- load_in_8bit=False,
31
  torch_dtype=torch.float16,
32
  device_map="auto",
33
  )
@@ -37,6 +37,7 @@ if device == "cuda":
37
  elif device == "mps":
38
  model = LlamaForCausalLM.from_pretrained(
39
  BASE_MODEL,
 
40
  device_map={"": device},
41
  torch_dtype=torch.float16,
42
  )
 
27
  if device == "cuda":
28
  model = LlamaForCausalLM.from_pretrained(
29
  BASE_MODEL,
30
+ load_in_8bit=True,
31
  torch_dtype=torch.float16,
32
  device_map="auto",
33
  )
 
37
  elif device == "mps":
38
  model = LlamaForCausalLM.from_pretrained(
39
  BASE_MODEL,
40
+ load_in_8bit=True,
41
  device_map={"": device},
42
  torch_dtype=torch.float16,
43
  )