BK-Lee commited on
Commit
2b9d94f
1 Parent(s): cbb062a
Files changed (1) hide show
  1. meteor/load_meteor.py +2 -2
meteor/load_meteor.py CHANGED
@@ -17,7 +17,7 @@ def load_meteor(link, bits):
17
  huggingface_config.update(dict(
18
  torch_dtype=torch.float16,
19
  low_cpu_mem_usage=True,
20
- # attn_implementation="flash_attention_2",
21
  quantization_config=BitsAndBytesConfig(
22
  load_in_4bit=bits == 4,
23
  load_in_8bit=bits == 8,
@@ -33,7 +33,7 @@ def load_meteor(link, bits):
33
  huggingface_config.update(dict(
34
  torch_dtype=torch.float16,
35
  low_cpu_mem_usage=True,
36
- # attn_implementation="flash_attention_2",
37
  ))
38
 
39
  # loading backbone model
 
17
  huggingface_config.update(dict(
18
  torch_dtype=torch.float16,
19
  low_cpu_mem_usage=True,
20
+ attn_implementation="flash_attention_2",
21
  quantization_config=BitsAndBytesConfig(
22
  load_in_4bit=bits == 4,
23
  load_in_8bit=bits == 8,
 
33
  huggingface_config.update(dict(
34
  torch_dtype=torch.float16,
35
  low_cpu_mem_usage=True,
36
+ attn_implementation="flash_attention_2",
37
  ))
38
 
39
  # loading backbone model