   input = THCudaTensor_newContiguous(input);
   THCudaTensor_resizeAs(output, input);
 
-  thrust::device_ptr<float> output_data(THCudaTensor_data(output));
-  thrust::device_ptr<float> input_data(THCudaTensor_data(input));
-
-  float logsum = thrust::reduce(input_data, input_data+size, (float)(-THInf), logaddforward_functor());
-  thrust::transform(input_data, input_data+size, output_data, addvalue_functor(-logsum));
+  if(input->nDimension == 1)
+  {
+    dim3 blocks(1);
+    dim3 threads(LOGSOFTMAX_THREADS);
+    cunn_LogSoftMax_forward_kernel<<<blocks,threads>>>(THCudaTensor_data(output), THCudaTensor_data(input), 1, input->size[0]);
+  }
+  else if(input->nDimension == 2)
+  {
+    dim3 blocks(input->size[0]);
+    dim3 threads(LOGSOFTMAX_THREADS);
+    cunn_LogSoftMax_forward_kernel<<<blocks,threads>>>(THCudaTensor_data(output), THCudaTensor_data(input), input->size[0], input->size[1]);
+  }
+  else
+    THError("vector or matrix expected");
+
+  cudaError errcode = cudaGetLastError();
+  if(errcode != cudaSuccess)
+    THError(cudaGetErrorString(errcode));
 
   THCudaTensor_free(input);
   return 1;

#define MINUS_LOG_THRESHOLD -18.42

struct logaddforward_functor
{
  __host__ __device__ float operator()(const float& log_a_, const float& log_b_) const
  {
    float minusdif;
    float log_a = log_a_;
    float log_b = log_b_;

    if (log_a < log_b)
    {
      log_a = log_b_;
      log_b = log_a_;
    }

    minusdif = log_b - log_a;

    if (minusdif < MINUS_LOG_THRESHOLD)
      return log_a;
    else
      return log_a + log1p(exp(minusdif));
  }
};

struct addvalue_functor
{
  const float value;

  addvalue_functor(float value_) : value(value_) {}

    __host__ __device__ float operator()(const float& x) const
  {
    return (x+value);
  }
};

static int cunn_LogSoftMax_forward(lua_State *L)
{


  THCudaTensor_resizeAs(gradInput, output);

  if(gradInput->nDimension == 1)
  {
    dim3 blocks(1);
    dim3 threads(LOGSOFTMAX_THREADS);

    cunn_LogSoftMax_backward_kernel<<<blocks,threads>>>(THCudaTensor_data(gradInput),
                                                        THCudaTensor_data(output),
                                                        THCudaTensor_data(gradOutput),
                                                        1, gradInput->size[0]);
  }
  else if(gradInput->nDimension == 2)
  {
    dim3 blocks(gradInput->size[0]);
    dim3 threads(LOGSOFTMAX_THREADS);

    cunn_LogSoftMax_backward_kernel<<<blocks,threads>>>(THCudaTensor_data(gradInput),
                                                        THCudaTensor_data(output),
                                                        THCudaTensor_data(gradOutput),
                                                        gradInput->size[0], gradInput->size[1]);
  }
  else
    THError("vector or matrix expected");

  cudaError errcode = cudaGetLastError();
  if(errcode != cudaSuccess)
    THError(cudaGetErrorString(errcode));

  THCudaTensor_free(gradOutput);
  THCudaTensor_free(output);
