From 551be9d93912e1a6dd69b6f1b254a71de8e9fd38 Mon Sep 17 00:00:00 2001
From: WuxinrongY <53141838+WuxinrongY@users.noreply.github.com>
Date: Tue, 26 Mar 2024 13:30:00 +0800
Subject: [PATCH] update gen_wts.py (#1467)

---
 yolov9/gen_wts.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/yolov9/gen_wts.py b/yolov9/gen_wts.py
index a507e1f..c90154b 100644
--- a/yolov9/gen_wts.py
+++ b/yolov9/gen_wts.py
@@ -32,7 +32,7 @@ print(f'Generating .wts for {m_type} model')
 print(f'Loading {pt_file}')
 device = select_device('cpu')
 model = torch.load(pt_file, map_location=device)  # Load FP32 weights
-model.model.float()
+model = model["model"].float()
 
 if m_type in ['detect', 'seg']:
     # update anchor_grid info