- Fix to GPU

test=develop

- Fix to GPU
test=develop
be80bb4f · Jacek Czaja · 513bb6c1 · be80bb4f
隐藏空白更改
内联并排

Showing with 3 addition and 1 deletion

paddle/fluid/operators/softmax_op.h paddle/fluid/operators/softmax_op.h +3 -1

未找到文件。
--- a/paddle/fluid/operators/softmax_op.h
+++ b/paddle/fluid/operators/softmax_op.h
@@ -36,7 +36,9 @@ class SoftmaxKernel : public framework::OpKernel<T> {
    Tensor Out_2d = framework::ReshapeToMatrix(*Out, rank - 1);

 #ifdef PADDLE_ON_INFERENCE
-    math::SoftmaxFunctor<DeviceContext, T, true>()(
+    math::SoftmaxFunctor<
+        DeviceContext, T,
+        std::is_same<DeviceContext, platform::CPUDeviceContext>::value>()(
        context.template device_context<DeviceContext>(), &X_2d, &Out_2d);
 #else
    math::SoftmaxFunctor<DeviceContext, T, false>()(