diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py index ed2409a2e2d4561141ffcc434b67139b304dcefb..b05261b4c66e8734c346c0566a4ec74e6257da9b 100644 --- a/python/paddle/fluid/layers/nn.py +++ b/python/paddle/fluid/layers/nn.py @@ -5712,7 +5712,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)`. @@ -5732,9 +5732,9 @@ def squeeze(input, axes, name=None): """ helper = LayerHelper("squeeze", **locals()) - check_type_and_dtype(input, 'input', Variable, - ['float32', 'float64', 'int8', 'int32', 'int64'], - 'squeeze') + check_type_and_dtype( + input, 'input', Variable, + ['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 74ebf2d208eeaca739d77d11763d016f868d7643..cb3b13cf160bb76a4369b12f2a449e4a6112befe 100644 --- a/python/paddle/fluid/layers/tensor.py +++ b/python/paddle/fluid/layers/tensor.py @@ -448,7 +448,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. @@ -469,9 +469,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)