Add metadata when converting .bin to .safetensors

TMK04 · Jan 10, 2024 · 699d127 · 699d127
1 parent 7d37b50
commit 699d127
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/util/convert_safetensors.py b/util/convert_safetensors.py
@@ -1,5 +1,5 @@
 import torch
-import argparse, os, glob, sys
+import argparse, os
 from safetensors.torch import save_file
 
 parser = argparse.ArgumentParser(description="Convert .bin/.pt files to .safetensors")
@@ -12,4 +12,4 @@
 
     out_file = os.path.splitext(file)[0] + ".safetensors"
     print(f" -- Saving {out_file}...")
-    save_file(state_dict, out_file)
+    save_file(state_dict, out_file, metadata = {'format': 'pt'})