我使用自己的数据集训练，类别为2，转化ONNX时报错

转化命令python3 tools/export_onnx.py -n yolox-s

RuntimeError: Error(s) in loading state_dict for YOLOX:
size mismatch for head.cls_preds.0.weight: copying a param with shape torch.Size([2, 128, 1, 1]) from checkpoint, the shape in current model is torch.Size([80, 128, 1, 1]).
size mismatch for head.cls_preds.0.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([80]).
size mismatch for head.cls_preds.1.weight: copying a param with shape torch.Size([2, 128, 1, 1]) from checkpoint, the shape in current model is torch.Size([80, 128, 1, 1]).
size mismatch for head.cls_preds.1.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([80]).
size mismatch for head.cls_preds.2.weight: copying a param with shape torch.Size([2, 128, 1, 1]) from checkpoint, the shape in current model is torch.Size([80, 128, 1, 1]).
size mismatch for head.cls_preds.2.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([80]).
怀疑最后输出类别没有修改为2
应该在哪里修改呢？

Start with tools/export_onnx.py and the yolox-s export command, then trace how the checkpoint's two-class shape is compared with the model's 80-class head. Done means export completes with a model configured for two classes; the issue provides no test path.