Update README.md

README.md

@@ -9,7 +9,6 @@ Rehosted from https://hf.co/mayrajeo/marine-vessel-yolo using the following code
 import os
 import hashlib
 
-import torch
 import ultralytics
 
 url = "https://hf.co/mayrajeo/marine-vessel-yolo/resolve/main/yolo11s_tci.pt"
@@ -27,7 +26,8 @@ print(model.args)
 
 # save and compute hash
 filename = "yolo11s_sentinel2_rgb_marine_vessel_detection.pt"
-torch.save(model.state_dict(), filename)
+model.ckpt = model.state_dict()
+model.save(filename)
 md5 = hashlib.md5(open(filename, "rb").read()).hexdigest()[:8]
 os.rename(filename, filename.replace(".pt", f"-{md5}.pt"))