diff --git a/test_load_quantized.py b/test_load_quantized.py
index 3203d42..11ff323 100644
--- a/test_load_quantized.py
+++ b/test_load_quantized.py
@@ -1,18 +1,26 @@
 #!/usr/bin/env python3
 """Test loading quantized shard to check for MISMATCH."""
 
+import gc
 import torch
+import shutil
+from pathlib import Path
 from transformers import AutoModelForCausalLM, AutoConfig
 
 
 def test_load():
-    print("Loading original model config...")
-    config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
+    # Copy config.json from original model
+    test_dir = Path("/data/models/test_quantize")
+    config_src = Path("/data/models/Ornith-1.0-35B") / "config.json"
+    config_dst = test_dir / "config.json"
+    if not config_dst.exists():
+        shutil.copy2(config_src, config_dst)
+        print(f"Copied config.json to {test_dir}")
     
     print("\nLoading quantized test shard...")
     try:
         model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/test_quantize",
+            str(test_dir),
             device_map="cpu",
             torch_dtype=torch.float16,
             trust_remote_code=True,