yolov5s.pt微调后在生成.wts报错，改了数据类别，但类别总数不变。麻烦大大帮忙看看~谢谢！！new.pt是我微调后的模型

[USTBZcb]

jetson@yahboom:~/yolov5$ python3 gen_wts.py -w new.pt
Traceback (most recent call last):
File "gen_wts.py", line 33, in
model = torch.load(pt_file, map_location=device)
File "/home/jetson/.local/lib/python3.6/site-packages/torch/serialization.py", line 592, in load
return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
File "/home/jetson/.local/lib/python3.6/site-packages/torch/serialization.py", line 851, in _load
result = unpickler.load()
AttributeError: Can't get attribute 'DetectionModel' on <module 'models.yolo' from '/home/jetson/yolov5/models/yolo.py'>

[wang-xinyu]

gen_wts在训练机器上导出即可