how to convert rwkv-5 model?

[xinyinan9527]

i used the convert.py to convert the rwkv-5 model ,when i used the test.py to test if it was successful,i failed.

[?] Select model: 
 > RWKV_24_2048_32_15.onnx

['AzureExecutionProvider', 'CPUExecutionProvider']
[?] Select execution providers(use space bar to select checkboxes): 
 > [ ] AzureExecutionProvider
   [ ] CPUExecutionProvider

[]
2024-05-18 15:30:52.673417593 [W:onnxruntime:, graph.cc:3593 CleanUnusedInitializersAndNodeArgs] Removing initializer 'PreTrainedTensor_4'. It is not used by any node and should be removed from the model.
2024-05-18 15:30:52.673432693 [W:onnxruntime:, graph.cc:3593 CleanUnusedInitializersAndNodeArgs] Removing initializer 'PreTrainedTensor_1'. It is not used by any node and should be removed from the model.
2024-05-18 15:30:52.673435491 [W:onnxruntime:, graph.cc:3593 CleanUnusedInitializersAndNodeArgs] Removing initializer 'PreTrainedTensor_0'. It is not used by any node and should be removed from the model.
tensor(float) 2048 24
(48, 2048)
(24, 40, 64, 64)
/.local/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
  0%|                                                    | 0/19 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "rwkv-onnx/ds/../test.py", line 120, in <module>
    logits, state, state2 = model.forward(token,state, state2)
  File "rwkv-onnx/ds/../test.py", line 66, in forward
    outputs = sess.run(output_names, inputs)
  File "/.conda/envs/py10/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 220, in run
    return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: instatewkv9 for the following indices
 index: 0 Got: 40 Expected: 32
 Please fix either the inputs/outputs or the model.

[harrisonvanderbyl]

https://github.com/RWKV/rwkv-onnx/blob/main/test.py#L75C47-L75C49
looks like I accidentally hard coded the head size to 40,
just change that 40 on this line to 32 to test your model