I find we set handle type be float32, however, for some ops, like ArgMax, we need to make output’s dtype is int32. Can we have one mechanism to set it?
otherwise we will get error here:
Expr cond = (TVMArrayGet(UInt(8), handle, intrinsic::kArrTypeCode) ==
** UIntImm::make(UInt(8), dtype.code())** &&
TVMArrayGet(UInt(8), handle, intrinsic::kArrTypeBits) ==
UIntImm::make(UInt(8), dtype.bits()) &&
TVMArrayGet(UInt(16), handle, intrinsic::kArrTypeLanes) ==
UIntImm::make(UInt(16), dtype.lanes()));
Because handle type is float32 when input is float32, we can not change it.
Thanks