device map fix

model_setup.py

@@ -4,8 +4,8 @@ import torch
 import gc
 
 from sentence_transformers import SentenceTransformer
-from transformers import AutoProcessor, Gemma3ForConditionalGeneration, BitsAndBytesConfig
-from accelerate import disk_offload
+from transformers import AutoProcessor, Gemma3ForConditionalGeneration
+# from accelerate import disk_offload
 from utils import clear_gpu_cache
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -15,24 +15,15 @@ embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
 # Gemma3 quantization config
 model_name = "google/gemma-3-4b-it"
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    # llm_int8_enable_fp32_cpu_offload=True  # Allow offloading    
-)
 
 # Load Gemma3
 model = Gemma3ForConditionalGeneration.from_pretrained(
     model_name,
     torch_dtype=torch.bfloat16,
-    device_map="cpu", # not "auto" since there is no GPU
-    quantization_config=bnb_config,
-    low_cpu_mem_usage=True,
-    # attn_implementation="sdpa"
+    device_map="cpu", # Avoid meta errors
 )
-disk_offload(model=model, offload_dir="offload")
+# disk_offload(model=model, offload_dir="offload")
+model.to("cpu")
 model.eval()
 
 # Processor