Tao Wu commited on
Commit
9fbbd17
·
1 Parent(s): 5d61362
Files changed (1) hide show
  1. app/embedding_setup.py +1 -1
app/embedding_setup.py CHANGED
@@ -45,7 +45,7 @@ first_id = tokenizer.convert_tokens_to_ids(first_token)
45
  second_id = tokenizer.convert_tokens_to_ids(second_token)
46
  model = AutoModelForCausalLM.from_pretrained(
47
  LLM_MODEL,
48
- load_in_8bit=True,
49
  torch_dtype=torch.float16,
50
  token=hf_auth,
51
  )
 
45
  second_id = tokenizer.convert_tokens_to_ids(second_token)
46
  model = AutoModelForCausalLM.from_pretrained(
47
  LLM_MODEL,
48
+ load_in_4bit=True,
49
  torch_dtype=torch.float16,
50
  token=hf_auth,
51
  )