From be411d69f45defdee8a658d02877a4dfbb338507 Mon Sep 17 00:00:00 2001
From: root
 <root@notebook-b086b94b-fa1a-11ef-ae8c-564dca19abc0-0.notebook-b086b94b-fa1a-11ef-ae8c-564dca19abc0.hpc-ai-cloud.svc.cluster.local>
Date: Thu, 6 Mar 2025 14:25:47 +0800
Subject: [PATCH] Fix: add metadata to bf16 safetensors for loading using
 transformers

---
 inference/fp8_cast_bf16.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/inference/fp8_cast_bf16.py b/inference/fp8_cast_bf16.py
index 4037342..ffb5f23 100644
--- a/inference/fp8_cast_bf16.py
+++ b/inference/fp8_cast_bf16.py
@@ -85,7 +85,7 @@ def main(fp8_path, bf16_path):
                 new_state_dict[weight_name] = weight
                 
         new_safetensor_file = os.path.join(bf16_path, file_name)
-        save_file(new_state_dict, new_safetensor_file)
+        save_file(new_state_dict, new_safetensor_file, metadata={"format": "pt"})
         
         # Memory management: keep only the 2 most recently used files
         if len(loaded_files) > 2: