softmax_op.cu 265 字节
Newer Older
Q
qijun 已提交
1
#define EIGEN_USE_GPU
2 3
#include "paddle/framework/op_registry.h"
#include "paddle/operators/softmax_op.h"
4

5
REGISTER_OP_GPU_KERNEL(softmax, ops::SoftmaxKernel<ops::GPUPlace, float>);
Q
Qiao Longfei 已提交
6
REGISTER_OP_GPU_KERNEL(softmax_grad, ops::SoftmaxGradKernel<ops::GPUPlace, float>);