From f438318583badd5e51746fdee73503393f4cfd57 Mon Sep 17 00:00:00 2001 From: Xun Deng Date: Thu, 27 Aug 2020 15:43:22 -0400 Subject: [PATCH] fix common_dtype in distribution utils --- mindspore/nn/probability/distribution/_utils/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mindspore/nn/probability/distribution/_utils/utils.py b/mindspore/nn/probability/distribution/_utils/utils.py index 39c36a5042..729c42634d 100644 --- a/mindspore/nn/probability/distribution/_utils/utils.py +++ b/mindspore/nn/probability/distribution/_utils/utils.py @@ -379,8 +379,12 @@ def common_dtype(arg_a, name_a, arg_b, name_b, hint_type): if hasattr(arg_a, 'dtype') and hasattr(arg_b, 'dtype'): if isinstance(arg_a, np.ndarray): a_dtype = mstype.pytype_to_dtype(arg_a.dtype) - if isinstance(arg_a, np.ndarray): + else: + a_dtype = arg_a.dtype + if isinstance(arg_b, np.ndarray): b_dtype = mstype.pytype_to_dtype(arg_b.dtype) + else: + b_dtype = arg_b.dtype if a_dtype != b_dtype: raise TypeError(f"{name_a} and {name_b} should have the same dtype.") int_type = mstype.int_type + mstype.uint_type