shauray commited on
Commit
136fb02
1 Parent(s): b95c9fa

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +0 -4
model.py CHANGED
@@ -6,19 +6,15 @@ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIt
6
 
7
  model_id = 'abacaj/starcoderbase-1b-sft'
8
 
9
- config = AutoConfig.from_pretrained(model_id)
10
- config.pretraining_tp = 1
11
  if torch.cuda.is_available():
12
  model = AutoModelForCausalLM.from_pretrained(
13
  model_id,
14
- config=config,
15
  torch_dtype=torch.float16,
16
  device_map='cuda',
17
  )
18
  else:
19
  model = AutoModelForCausalLM.from_pretrained(
20
  model_id,
21
- config=config,
22
  torch_dtype=torch.float32,
23
  )
24
  tokenizer = AutoTokenizer.from_pretrained(model_id)
 
6
 
7
  model_id = 'abacaj/starcoderbase-1b-sft'
8
 
 
 
9
  if torch.cuda.is_available():
10
  model = AutoModelForCausalLM.from_pretrained(
11
  model_id,
 
12
  torch_dtype=torch.float16,
13
  device_map='cuda',
14
  )
15
  else:
16
  model = AutoModelForCausalLM.from_pretrained(
17
  model_id,
 
18
  torch_dtype=torch.float32,
19
  )
20
  tokenizer = AutoTokenizer.from_pretrained(model_id)