Error saving final weights

[zakizuberiaria]

Getting an error in line 371 of train.py when saving final weights. It seems the arguments are incorrect. I believe the line should be updated totrainer.export_model(model, optimizer, name="final") to work.

[zakizuberiaria]

Additionally, the step is not saved for the final weights causing an error in line 136 of predict when running inference. However, step is not actually used in predict except for display purposes. Step needs to be saved in train.py line 371 or not loaded in predict.py line 136.