diff --git a/train.py b/train.py
index fc5bfb8..13f344e 100644
--- a/train.py
+++ b/train.py
@@ -35,11 +35,36 @@ def train(config_path):
     
     # Check if model is quantized (has model-XXXXX-of-XXXXX files)
     import glob
+    import json
     safetensor_files = glob.glob(f"{config['base_model']}/*.safetensors")
-    is_quantized = any("of-" in f for f in safetensor_files)
+    is_quantized = any("of-" in Path(f).name for f in safetensor_files)
     
     if is_quantized:
-        print(f"\n[INFO] Loading pre-quantized BnB 4-bit model (no re-quantization)...")
+        # Create proper index file from safetensors metadata
+        index_file = Path(config["base_model"]) / "model.safetensors.index.json"
+        if not index_file.exists():
+            print(f"\n[INFO] Creating model index file from safetensors metadata...")
+            shards = sorted([Path(f).name for f in safetensor_files if "of-" in Path(f).name])
+            
+            # Load metadata from first shard to get weight names
+            from safetensors import safe_open
+            weight_map = {}
+            for shard_name in shards:
+                shard_path = Path(config["base_model"]) / shard_name
+                with safe_open(str(shard_path), framework="pt", device="cpu") as f:
+                    for key in f.keys():
+                        weight_map[key] = shard_name
+            
+            # Create index
+            index = {
+                "metadata": {"total_size": sum((Path(config["base_model"]) / s).stat().st_size for s in shards)},
+                "weight_map": weight_map
+            }
+            with open(index_file, 'w') as f:
+                json.dump(index, f)
+            print(f"✓ Created index: {index_file} ({len(weight_map)} weights mapped)")
+        
+        print(f"\n[INFO] Loading pre-quantized BnB 4-bit model...")
         model = AutoModelForCausalLM.from_pretrained(
             config["base_model"],
             device_map="cpu",