updated quantization config

app.py

@@ -1,15 +1,16 @@
 import gradio as gr
 import torch
 from peft import PeftModel, PeftConfig
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 
 #Loading model
+quantization_config = BitsAndBytesConfig(load_in_4bit=True)
+
 model_path = "parsanna17/finetune_starcoder2_with_R_data"
 checkpoint = "bigcode/starcoder2-3b"
-device = "cuda" if torch.cuda.is_available() else "cpu"
 config = PeftConfig.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(checkpoint , device_map=device, torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained(checkpoint, quantization_config=quantization_config)
 model = PeftModel.from_pretrained(model, model_path).to(device)
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)