diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py index 0cbcbf75171e9c9ca6506c69e9bf25b1f7d82fa1..f519ab9bb7ccad360d8ae6fc128764d0053b387a 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 cd93eca35ffc7af238d8e3c6fb5ec4b597e1d11b..0bbb57c97c17305efc3e6fe70e99d1b973e447c4 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)