Browse Source

!8599 Fix bug of Unique infershape with data type int64

From: @jojobugfree
Reviewed-by: @chujinjin,@jjfeing
Signed-off-by: @jjfeing
tags/v1.1.0
mindspore-ci-bot Gitee 5 years ago
parent
commit
23d34bc34f
1 changed files with 3 additions and 1 deletions
  1. +3
    -1
      mindspore/core/abstract/prim_arrays.cc

+ 3
- 1
mindspore/core/abstract/prim_arrays.cc 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());


Loading…
Cancel
Save