fix: use torch.load for quantized shards index creation

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 = {