teknium commited on
Commit
9ae7036
1 Parent(s): 93f2ad1

Update transformers_inference.py

Browse files
Files changed (1) hide show
  1. transformers_inference.py +1 -1
transformers_inference.py CHANGED
@@ -10,7 +10,7 @@ tokenizer = LlamaTokenizer.from_pretrained('teknium/OpenHermes-2.5-Mistral-7B',
10
  model = MistralForCausalLM.from_pretrained(
11
  "teknium/OpenHermes-2.5-Mistral-7B",
12
  torch_dtype=torch.float16,
13
- device_map=device_map="auto",#{'': 'cuda:0'},
14
  load_in_8bit=False,
15
  load_in_4bit=True,
16
  use_flash_attention_2=True
 
10
  model = MistralForCausalLM.from_pretrained(
11
  "teknium/OpenHermes-2.5-Mistral-7B",
12
  torch_dtype=torch.float16,
13
+ device_map="auto",#{'': 'cuda:0'},
14
  load_in_8bit=False,
15
  load_in_4bit=True,
16
  use_flash_attention_2=True