update gen_wts.py (#1467)

2024-03-26 13:30:00 +08:00 · 2024-03-26 13:30:00 +08:00 · 551be9d939
commit 551be9d939
parent f12bbe0359
1 changed files with 1 additions and 1 deletions
--- 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