Onnx_export.py outputs size mismatch for classification_headers.0.weight / bias errors

I am trying to retrain my own detection model on a Jetson Xavier NX, I have used camera-capture to grab some images and I can run train_ssd.py without any problems however when I then run onnx_export.py I get a string of errors like
Traceback (most recent call last):
File “onnx_export.py”, line 86, in
net.load(args.input)
File “/home/charles/jetson-inference/python/training/detection/ssd/vision/ssd/ssd.py”, line 135, in load
self.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage))
File “/home/charles/.local/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1045, in load_state_dict
self.class.name, “\n\t”.join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for SSD:
size mismatch for classification_headers.0.weight: copying a param with shape torch.Size([30, 512, 3, 3]) from checkpoint, the shape in current model is torch.Size([24, 512, 3, 3]).
size mismatch for classification_headers.0.bias: copying a param with shape torch.Size([30]) from checkpoint, the shape in current model is torch.Size([24]).

etc

can someone tell me where this issue may be coming from?
The input image size I am using are 1024x1024x3

  • Charles

Problem solved see here https://github.com/dusty-nv/jetson-inference/issues/820