shauray commited on
Commit
00dafd6
1 Parent(s): 2411fe0

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +4 -4
model.py CHANGED
@@ -10,14 +10,14 @@ if torch.cuda.is_available():
10
  model = AutoModelForCausalLM.from_pretrained(
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
  )
 
10
  model = AutoModelForCausalLM.from_pretrained(
11
  model_id,
12
  torch_dtype=torch.float16,
13
+ device_map='auto',
14
 
15
+ )
16
  else:
17
  model = AutoModelForCausalLM.from_pretrained(
18
  model_id,
19
+ torch_dtype=torch.float32,
20
+ #load_in_4bit=True,
21
  device_map='auto',
22
  use_safetensors=False,
23
  )