!8599 Fix bug of Unique infershape with data type int64

From: @jojobugfree
Reviewed-by: @chujinjin,@jjfeing
Signed-off-by: @jjfeing
This commit is contained in:
mindspore-ci-bot 2020-11-16 09:06:01 +08:00 committed by Gitee
commit 23d34bc34f
1 changed files with 3 additions and 1 deletions

View File

@ -170,7 +170,9 @@ AbstractBasePtr InferImplUnique(const AnalysisEnginePtr &, const PrimitivePtr &p
std::make_shared<AbstractTensor>(input->element(), std::make_shared<Shape>(ids_shape, min_shape, max_shape));
// Currently we choose the same data type as input for the idx.
TypePtr ids_idx_type = kInt32;
if (input->element() != nullptr && input->element()->GetTypeTrack() == kInt64) {
MS_EXCEPTION_IF_NULL(input->element());
MS_EXCEPTION_IF_NULL(input->element()->GetTypeTrack());
if (input->element()->GetTypeTrack()->type_id() == TypeId::kNumberTypeInt64) {
ids_idx_type = kInt64;
}
auto ids_idx = std::make_shared<AbstractTensor>(ids_idx_type, shape->shape());