shauray commited on
Commit
2411fe0
1 Parent(s): 05004d5

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +5 -1
model.py CHANGED
@@ -11,11 +11,15 @@ if torch.cuda.is_available():
11
  model_id,
12
  torch_dtype=torch.float16,
13
  device_map='cuda',
 
14
  ).to("cuda")
15
  else:
16
  model = AutoModelForCausalLM.from_pretrained(
17
  model_id,
18
- torch_dtype=torch.float32,
 
 
 
19
  )
20
  tokenizer = AutoTokenizer.from_pretrained(model_id)
21
 
 
11
  model_id,
12
  torch_dtype=torch.float16,
13
  device_map='cuda',
14
+
15
  ).to("cuda")
16
  else:
17
  model = AutoModelForCausalLM.from_pretrained(
18
  model_id,
19
+ torch_dtype=torch.float16,
20
+ load_in_4bit=True,
21
+ device_map='auto',
22
+ use_safetensors=False,
23
  )
24
  tokenizer = AutoTokenizer.from_pretrained(model_id)
25