type of return operand 0 ('!torch.vtensor<[?,384],f32>') doesn't match function result type ('!torch.vtensor<[1,384],f32>') in function @torch-jit-export

[pdhirajkumarprasad]

For the given IR

module {
  func.func @"torch-jit-export"(%arg4: !torch.vtensor<[?,384,1],f32>) -> (!torch.vtensor<[1,384],f32>) attributes {torch.onnx_meta.ir_version = 6 : si64, torch.onnx_meta.opset_version = 21 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "1.8"} {
    %1 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<2> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64> 
    %2 = torch.operator "onnx.Squeeze"(%arg4, %1) : (!torch.vtensor<[?,384,1],f32>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[1,384],f32> 
    return %2 : !torch.vtensor<[1,384],f32>

  }
}

getting error as

error: type of return operand 0 ('!torch.vtensor<[?,384],f32>') doesn't match function result type ('!torch.vtensor<[1,384],f32>') in function @torch-jit-export

command :

iree-compile --iree-hal-target-backends=llvm-cpu -o abc.vmfb model.torch_onnx.mlir --iree-llvmcpu-target-cpu=host