version 1.5

2024-11-26 13:04:51 -08:00
parent 7021ad8917
commit 856b539e54
3 changed files with 2 additions and 5 deletions
--- a/weights/convert_safetensor_to_pt.py
+++ b/weights/convert_safetensor_to_pt.py
@@ -16,7 +16,7 @@ if args.version == 'v1':
    torch.save({'model':model}, 'weights/icon_detect/best.pt')
 elif args.version == 'v1_5':
    print("Converting v1_5")
-    tensor_dict = load_file("weights/icon_detect_v1_5/model.safetensors")
+    tensor_dict = torch.load("weights/icon_detect_v1_5/model.safetensors")
    model = DetectionModel('weights/icon_detect_v1_5/model.yaml')
    model.load_state_dict(tensor_dict)
    save_dict = {'model':model}