From 1ebad2365ce4c7041ed721ab6ab90bce41284810 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 03:58:24 -0400
Subject: [PATCH] fix: use torch.load for quantized shards index creation

---
 train.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/train.py b/train.py
index d746cd7..23e7483 100644
--- a/train.py
+++ b/train.py
@@ -36,7 +36,6 @@ def train(config_path):
     # Create index file for sharded model
     import glob as glob_mod
     import json as json_mod
-    from safetensors import safe_open
     
     safetensor_files = glob_mod.glob(f"{config['base_model']}/*.safetensors")
     shards = sorted([Path(f).name for f in safetensor_files if "of-" in Path(f).name])
@@ -47,8 +46,9 @@ def train(config_path):
         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():
+            ckpt = torch.load(str(shard_path), map_location="cpu", weights_only=False)
+            for key in ckpt.keys():
+                if isinstance(ckpt[key], torch.Tensor):
                     weight_map[key] = shard_name
         
         index = {