softmax_op.cu 288 字节
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>);
L
liaogang 已提交
6 7
REGISTER_OP_GPU_KERNEL(softmax_grad,
                       ops::SoftmaxGradKernel<ops::GPUPlace, float>);