Update app.py

app.py

@@ -77,7 +77,8 @@ def initialize_model():
         load_in_8bit=True,
         bnb_8bit_compute_dtype=torch.bfloat16,
         bnb_8bit_quant_type="nf4",
-        bnb_8bit_use_double_quant=True
+        bnb_8bit_use_double_quant=True,
+        llm_int8_enable_fp32_cpu_offload=True
     )
 
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID , trust_remote_code=True)
@@ -86,7 +87,7 @@ def initialize_model():
 
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        torch_dtype=torch.float16,
+        torch_dtype="auto",
         device_map="cuda",
         # attn_implementation="flash_attention_2",
         trust_remote_code=True,