Key error when folding_constants

[JaclebusIII]

Here is a python script I used to convert an onnx model to rknn:

import numpy as np
import cv2
from rknn.api import RKNN
import os

ONNX_MODEL = '/airockchip/jack_sandbox/ast-finetuned-audioset-10-10-0.4593.onnx'
RKNN_MODEL = '/airockchip/jack_sandbox/ast-finetuned-audioset-10-10-0.4593.rknn'


if __name__ == "__main__":
    # Create RKNN object
    rknn = RKNN(verbose=True)

    # Pre-process config
    print("--> Config model")
    rknn.config( target_platform="rk3566", dynamic_input=None)
    print("done")

    # Load model
    print("--> Loading model")
    ret = rknn.load_onnx(model=ONNX_MODEL, inputs=['input_values'], outputs=["logits"], input_size_list=[[1, 1024, 128]])
    if ret != 0:
        print("Load model failed!")
        exit(ret)

    ret = rknn.build(do_quantization=False, rknn_batch_size=1)
    if ret != 0:
        print("Build model failed!")
        exit(ret)

    # Export rknn model
    print("--> Export RKNN model")
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print("Export RKNN model failed!")
        exit(ret)
    print("done")

I got the model from here:
https://huggingface.co/Xenova/ast-finetuned-audioset-10-10-0.4593/tree/main

When I run the script, I get this error:

D     convert_exmatmul_to_conv: remove node = ['/audio_spectrogram_transformer/encoder/layer.10/intermediate/dense/MatMul'], add node = ['/audio_spectrogram_transformer/encoder/layer.10/intermediate/dense/MatMul']
D     convert_exmatmul_to_conv: remove node = ['/audio_spectrogram_transformer/encoder/layer.10/output/dense/MatMul'], add node = ['/audio_spectrogram_transformer/encoder/layer.10/output/dense/MatMul']
D     convert_exmatmul_to_conv: remove node = ['/audio_spectrogram_transformer/encoder/layer.11/attention/attention/key/MatMul'], add node = ['/audio_spectrogram_transformer/encoder/layer.11/attention/attention/key/MatMul']
D     convert_exmatmul_to_conv: remove node = ['/audio_spectrogram_transformer/encoder/layer.11/attention/attention/value/MatMul'], add node = ['/audio_spectrogram_transformer/encoder/layer.11/attention/attention/value/MatMul']
D     convert_exmatmul_to_conv: remove node = ['/audio_spectrogram_transformer/encoder/layer.11/attention/attention/query/MatMul'], add node = ['/audio_spectrogram_transformer/encoder/layer.11/attention/attention/query/MatMul']
D     convert_exmatmul_to_conv: remove node = ['/audio_spectrogram_transformer/encoder/layer.11/attention/output/dense/MatMul'], add node = ['/audio_spectrogram_transformer/encoder/layer.11/attention/output/dense/MatMul']
D     convert_exmatmul_to_conv: remove node = ['/audio_spectrogram_transformer/encoder/layer.11/intermediate/dense/MatMul'], add node = ['/audio_spectrogram_transformer/encoder/layer.11/intermediate/dense/MatMul']
D     convert_exmatmul_to_conv: remove node = ['/audio_spectrogram_transformer/encoder/layer.11/output/dense/MatMul'], add node = ['/audio_spectrogram_transformer/encoder/layer.11/output/dense/MatMul']
D     convert_exmatmul_to_conv: remove node = ['/classifier/dense/Gemm'], add node = ['/classifier/dense/Gemm']
D     fold_constant ...
E build: Traceback (most recent call last):
  File "rknn/api/rknn_log.py", line 344, in rknn.api.rknn_log.error_catch_decorator.error_catch_wrapper
  File "rknn/api/rknn_base.py", line 1962, in rknn.api.rknn_base.RKNNBase.build
  File "rknn/api/graph_optimizer.py", line 2240, in rknn.api.graph_optimizer.GraphOptimizer.fuse_ops
  File "rknn/api/graph_optimizer.py", line 956, in rknn.api.graph_optimizer.GraphOptimizer.fold_constant
KeyError: '/audio_spectrogram_transformer/embeddings/patch_embeddings/Transpose_1_output_0'

W build: ===================== WARN(3) =====================
E rknn-toolkit2 version: 2.3.0
Traceback (most recent call last):
  File "rknn/api/rknn_log.py", line 344, in rknn.api.rknn_log.error_catch_decorator.error_catch_wrapper
  File "rknn/api/rknn_base.py", line 1962, in rknn.api.rknn_base.RKNNBase.build
  File "rknn/api/graph_optimizer.py", line 2240, in rknn.api.graph_optimizer.GraphOptimizer.fuse_ops
  File "rknn/api/graph_optimizer.py", line 956, in rknn.api.graph_optimizer.GraphOptimizer.fold_constant
KeyError: '/audio_spectrogram_transformer/embeddings/patch_embeddings/Transpose_1_output_0'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/airockchip/jack_sandbox/ast-model-conversion.py", line 26, in <module>
    ret = rknn.build(do_quantization=False, rknn_batch_size=1)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/rknn/api/rknn.py", line 192, in build
    return self.rknn_base.build(do_quantization=do_quantization, dataset=dataset, expand_batch_size=rknn_batch_size)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "rknn/api/rknn_log.py", line 349, in rknn.api.rknn_log.error_catch_decorator.error_catch_wrapper
  File "rknn/api/rknn_log.py", line 95, in rknn.api.rknn_log.RKNNLog.e
ValueError: Traceback (most recent call last):
  File "rknn/api/rknn_log.py", line 344, in rknn.api.rknn_log.error_catch_decorator.error_catch_wrapper
  File "rknn/api/rknn_base.py", line 1962, in rknn.api.rknn_base.RKNNBase.build
  File "rknn/api/graph_optimizer.py", line 2240, in rknn.api.graph_optimizer.GraphOptimizer.fuse_ops
  File "rknn/api/graph_optimizer.py", line 956, in rknn.api.graph_optimizer.GraphOptimizer.fold_constant
KeyError: '/audio_spectrogram_transformer/embeddings/patch_embeddings/Transpose_1_output_0'

I am running this in a docker container base on python3.11, installing this whl file:
rknn_toolkit2-2.3.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl