未验证 提交 a1350744 编写于 作者: Z Zhang Ting 提交者: GitHub

register fp16 kernel, test=develop (#25630)

上级 c4192a80
......@@ -14,8 +14,10 @@ limitations under the License. */
#include "paddle/fluid/operators/shape_op.h"
REGISTER_OP_CUDA_KERNEL(shape, paddle::operators::ShapeKernel<int>,
paddle::operators::ShapeKernel<int32_t>,
paddle::operators::ShapeKernel<int64_t>,
paddle::operators::ShapeKernel<float>,
paddle::operators::ShapeKernel<double>);
REGISTER_OP_CUDA_KERNEL(
shape, paddle::operators::ShapeKernel<int>,
paddle::operators::ShapeKernel<int32_t>,
paddle::operators::ShapeKernel<int64_t>,
paddle::operators::ShapeKernel<float>,
paddle::operators::ShapeKernel<double>,
paddle::operators::ShapeKernel<paddle::platform::float16>);
......@@ -11101,7 +11101,7 @@ def shape(input):
input.shape = [3, 2]
Args:
input (Variable): The input can be N-D Tensor or SelectedRows with data type float32, float64, int32, int64.
input (Variable): The input can be N-D Tensor or SelectedRows with data type float16, float32, float64, int32, int64.
If input variable is type of SelectedRows, returns the shape of it's inner tensor.
Returns:
......@@ -11124,8 +11124,9 @@ def shape(input):
res = exe.run(fluid.default_main_program(), feed={'x':img}, fetch_list=[output])
print(res) # [array([ 3, 100, 100], dtype=int32)]
"""
check_variable_and_dtype(input, 'input',
['float32', 'float64', 'int32', 'int64'], 'shape')
check_variable_and_dtype(
input, 'input', ['float16', 'float32', 'float64', 'int32', 'int64'],
'shape')
helper = LayerHelper('shape', **locals())
out = helper.create_variable_for_type_inference(dtype='int32')
helper.append_op(
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册