From b8eb52f6bd3d6346c08caf5c331bff644d60533f Mon Sep 17 00:00:00 2001 From: wangchaochaohu Date: Fri, 5 Jun 2020 20:14:48 +0800 Subject: [PATCH] fix the fp16 support of assgin and squeeze Op test=develop (#24939) --- python/paddle/fluid/layers/nn.py | 8 ++++---- python/paddle/fluid/layers/tensor.py | 9 +++++---- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py index 0cbcbf75171..f519ab9bb7c 100644 --- a/python/paddle/fluid/layers/nn.py +++ b/python/paddle/fluid/layers/nn.py @@ -7479,7 +7479,7 @@ def squeeze(input, axes, name=None): Out.shape = [1,3,5] Args: - input (Variable): The input Tensor. Support data type: float32, float64, int8, int32, int64. + input (Variable): The input Tensor. Support data type: float16, float32, float64, int8, int32, int64. axes (list): One integer or List of integers, indicating the dimensions to be squeezed. Axes range is :math:`[-rank(input), rank(input))`. If axes is negative, :math:`axes=axes+rank(input)`. @@ -7499,9 +7499,9 @@ def squeeze(input, axes, name=None): """ helper = LayerHelper("squeeze", **locals()) - check_variable_and_dtype(input, 'input', - ['float32', 'float64', 'int8', 'int32', 'int64'], - 'squeeze') + check_variable_and_dtype( + input, 'input', + ['float16', 'float32', 'float64', 'int8', 'int32', 'int64'], 'squeeze') check_type(axes, 'axes', list, 'squeeze') out = helper.create_variable_for_type_inference(dtype=input.dtype) x_shape = helper.create_variable_for_type_inference(dtype=input.dtype) diff --git a/python/paddle/fluid/layers/tensor.py b/python/paddle/fluid/layers/tensor.py index cd93eca35ff..0bbb57c97c1 100644 --- a/python/paddle/fluid/layers/tensor.py +++ b/python/paddle/fluid/layers/tensor.py @@ -566,7 +566,7 @@ def assign(input, output=None): Parameters: input (Variable|numpy.ndarray): A tensor or numpy ndarray, its data type supports - float32, float64, int32 and int64. + float16, float32, float64, int32 and int64. output (Variable, optional): A tensor. If :attr:`output` is None, a new tensor will be created as :attr:`output`. Default: None. @@ -587,9 +587,10 @@ def assign(input, output=None): helper = LayerHelper('assign', **locals()) check_type(input, 'input', (Variable, numpy.ndarray), 'assign') if isinstance(input, Variable): - check_dtype(input.dtype, 'input', - ['float32', 'float64', 'int32', 'int64', 'bool'], 'assign', - '(When the type of input in assign is Variable.)') + check_dtype( + input.dtype, 'input', + ['float16', 'float32', 'float64', 'int32', 'int64', 'bool'], + 'assign', '(When the type of input in assign is Variable.)') if output is None: output = helper.create_variable_for_type_inference( dtype=input.dtype) -- GitLab