RWKV-6 模型转换报错 `invalid tensor malloc size, tensor name: , target: CPU, size: 0` (rknn-toolkit2 v2.1.0)

#!/usr/bin/env python
# coding: utf-8

from rknn.api import RKNN
from sys import exit
rknn = RKNN(verbose=True)

ONNX_MODEL="RWKV-x060-World-1B6-v2.1-20240328-ctx4096.onnx"
RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
DATASET="dataset.txt"
QUANTIZE=False

batch_size = 1

# pre-process config
print('--> Config model')
rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3)
print('done')

# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL, inputs=
                     ['/emb/Gather_output_0',
                      'input_state',
                      'scale_ratio'
                      ],
                      input_size_list=[
                          [batch_size, 2048],
                          [batch_size, 1584, 2048],
                          [1]
                        ],
                        outputs=[
                            'out',
                            'out_state',
                            '/blocks.0/att_group_norm/Reshape_output_0',
                            '/blocks.0/att_group_norm/InstanceNormalization_output_0'
                        ])
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')

# Build model
print('--> Building model')
ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
if ret != 0:
    print('Build model failed!')
    exit(ret)
print('done')

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

转换时报错如下

D RKNN: [20:37:25.716] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:25.717] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:25.717] finish initComputeZoneMapByStepsVector
W RKNN: [20:37:26.099] tensor name: scale_ratio_rs_tp_weight_0float16_2048_1_1_1 is already in IndependentMap
D RKNN: [20:37:26.100] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.100] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.100] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.100] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.100] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.100] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.107] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.107] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.107] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.107] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.109] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.109] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.109] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.109] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.109] finish initComputeZoneMapByStepsVector
D RKNN: [20:37:26.109] finish initComputeZoneMapByStepsVector
E RKNN: [20:37:26.109] invalid tensor malloc size, tensor name: , target: CPU, size: 0
fish: Job 1, 'python ./convert_rknn.py' terminated by signal SIGSEGV (Address boundary error)

2.0.0b17版本没有这个问题

airockchip / rknn-toolkit2

RWKV-6 模型转换报错 `invalid tensor malloc size, tensor name: , target: CPU, size: 0` (rknn-toolkit2 v2.1.0) #168