TVM compilation failed on SimpleViT

[yangxin0926]

I use this repo to export a SimpleViT onnx model. And I tried to compile the onnx model with relay but I got the following error.

Check failed: src_idx < ishape.size() (3 vs. 3) :

I thought this was raised by multihead for after changing heads to 1 the error does not occur.
But after I changed the heads to 1, I've got another error.

Check failed: (reporter->AssertEQ(xk, yk)) is false: BatchDot: shapes of x and y is inconsistent,  x shape=[1, 64, 1], y shape=[1, 64, 1024]

I think this error is raised by node /transformer/layers.0.1/net/net.0/ReduceMean because it's the last node that relay processed.
I don't know what to solve these problems. Here is my code generating onnx model.

from vit_pytorch import SimpleViT

image_size = 256
patch_size = 32
num_classes = 1000
dim = 1024
depth = 6
heads = 1
mlp_dim = 2048

model = SimpleViT(
        image_size = image_size,
        patch_size = patch_size,
        num_classes = num_classes,
        dim = dim,
        depth = depth,
        heads = heads,
        mlp_dim = mlp_dim
)

model.eval()

dummy_input = torch.randn(1, 3, image_size, image_size, requires_grad=True)

print(model)

print(model(dummy_input))

torch.onnx.export(
        model,
        dummy_input,
        "ViT.onnx",
        export_params = True,
        opset_version = 11,
        input_names = ['imageInput'],
        output_names = ['predOutput'],
        dynamic_axes = {'imageInput' : {0 : 'batch_size'}, 'predOutput' : {0 : 'batch_size'}})

And here is my code to compile the onnx model.

import onnx
import tvm
from tvm import te
import tvm.relay as relay

model = onnx.load("ViT.onnx")

shape_dict = {"imageInput" : [1, 3, 256, 256]}
mod, params = relay.frontend.from_onnx(model, shape_dict)