 
   input = THCudaTensor_newContiguous(state, input);
   THCudaTensor_resizeAs(state, output, input);
+  long batchSize, dim, stride;
 
   if(input->nDimension == 1)
   {
-    dim3 blocks(1);
-    dim3 threads(SOFTMAX_THREADS);
-    cunn_SoftMax_updateOutput_kernel<<<blocks,threads,
-      0, THCState_getCurrentStream(state)>>>(THCudaTensor_data(state, output),
-                                             THCudaTensor_data(state, input),
-                                             1, input->size[0]);
+    batchSize = 1;
+    dim = input->size[0];
+    stride = 1;
   }
   else if(input->nDimension == 2)
   {
-    dim3 blocks(input->size[0]);
-    dim3 threads(SOFTMAX_THREADS);
-    cunn_SoftMax_updateOutput_kernel<<<blocks,threads,
-      0, THCState_getCurrentStream(state)>>>(THCudaTensor_data(state, output),
-                                             THCudaTensor_data(state, input),
-                                             input->size[0], input->size[1]);
+    batchSize = input->size[0];
+    dim = input->size[1];
+    stride = 1;
+  }
+  else if(input->nDimension == 3)
+  {
+    batchSize = 1;
+    dim = input->size[0];
+    stride = input->size[1]*input->size[2];
+  }
+  else if(input->nDimension == 4)
+  {
+    batchSize = input->size[0];
+    dim = input->size[1];
+    stride = input->size[2]*input->size[3];
   }
   else
-    THError("vector or matrix expected");
+    THError("1D, 2D, 3D or 4D tensor expected");
+
+  dim3 blocks(batchSize, stride);
+  dim3 threads(SOFTMAX_THREADS);
+  cunn_SoftMax_updateOutput_kernel<<<blocks,threads,
+    0, THCState_getCurrentStream(state)>>>(THCudaTensor_data(state, output),
+                                           THCudaTensor_data(state, input),
+                                           batchSize, dim, stride);
 
   cudaError errcode = cudaGetLastError();
   if(errcode != cudaSuccess)

  gradOutput = THCudaTensor_newContiguous(state, gradOutput);

  THCudaTensor_resizeAs(state, gradInput, output);

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

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

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

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


  float sum_k = buffer[0];
  for (int i=i_start; i<i_end; i+=i_step)
    gradInput_k[i*stride] = output_k[i*stride] * (gradOutput_k[i*stride] - sum_k);
}

static int cunn_SoftMax_updateOutput(lua_State *L)
