Update model.py

model.py

@@ -6,24 +6,20 @@ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIt
 
 model_id = 'abacaj/starcoderbase-1b-sft'
 
+config = AutoConfig.from_pretrained(model_id)
+config.pretraining_tp = 1
 if torch.cuda.is_available():
-    config = AutoConfig.from_pretrained(model_id)
-    config.pretraining_tp = 1
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         config=config,
         torch_dtype=torch.float16,
-        load_in_4bit=True,
         device_map='cuda',
     )
 else:
-    config = AutoConfig.from_pretrained(model_id)
-    config.pretraining_tp = 1
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         config=config,
         torch_dtype=torch.float32,
-        load_in_4bit=True,
     )
 tokenizer = AutoTokenizer.from_pretrained(model_id)