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