           workspace_size,
           phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
 
-      auto* dx_data = dev_ctx->Alloc<T>(dx, dx->numel() * sizeof(T));
+      auto* dx_data = dev_ctx.Alloc<T>(dx, dx->numel() * sizeof(T));
       const auto* y_data = y->data<T>();
       const auto* dout_data = dout->data<T>();
       const auto* a_data = kXGradAIsDZ ? dout_data : y_data;

          dev_ctx.GetPlace(),
          workspace_size,
          phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
      auto* dy_data = dev_ctx->Alloc<T>(dy, dy->numel() * sizeof(T));
      const auto* dout_data = dout->data<T>();
      const auto* x_data = x->data<T>();
      const auto* a_data = kYGradAIsDZ ? dout_data : x_data;

      } else {
        reserve_space_size = phi::product(out->dims()) * sizeof(T);
      }
      dev_ctx.Alloc(reserve_space, out->type(), reserve_space_size);
      void* aux_data = reinterpret_cast<void*>(reserve_space->data<T>());

      PADDLE_ENFORCE_GPU_SUCCESS(
