     __syncthreads();
 
     // logadd?
-    float max_k = buffer[LOGSOFTMAX_THREADS];
-    buffer[threadIdx.x] = 0;
-    for (int i=i_start; i<i_end; i+=i_step)
-      buffer[threadIdx.x] += __expf(input_k[i]-max_k);
+    float max_k = buffer[SOFTMAXTREE_THREADS];
+    buffer[i_start] = 0;
+    for (int i=i_start; i<nOutput; i+=i_step)
+      buffer[i_start] += __expf(nodeInter[i]-max_k);
 
     __syncthreads();
 
     // reduce
-    if (threadIdx.x == 0)
+    if (i_start == 0)
     {
       float logsum_k = 0;
-      for (int i=0; i<blockDim.x; i++)
+      for (int i=0; i<nOutput; i++)
         logsum_k += buffer[i];
-      buffer[LOGSOFTMAX_THREADS] = max_k + __logf(logsum_k);
+      buffer[SOFTMAXTREE_THREADS] = max_k + __logf(logsum_k);
     }
 
     __syncthreads();
 
     // logsoftmax
-    float logsum_k = buffer[LOGSOFTMAX_THREADS];
-    for (int i=i_start; i<i_end; i+=i_step)
-      output_k[i] = input_k[i] - logsum_k;
-      
-    /* Narrow */
-    THTensor_(set)(nodeInter, nodeOutput);
-    THTensor_(narrow)(nodeOutput, nodeInter, 0, childIdx, 1); 
-    
-    /* CAddTable (without log, would have been CMulTable) */
-    narrowsum += THTensor_(get1d)(nodeOutput, 0);
+    float logsum_k = buffer[SOFTMAXTREE_THREADS];
+    for (int i=i_start; i<nOutput; i+=i_step)
+      nodeOutput[i] = nodeInter[i] - logsum_k;
+      
+    /* Narrow + CAddTable (without log, would have been CMulTable) */
+    narrowsum += nodeOutput[childIdx]
     
     n += nChildren;
     /* Break when root is reached */

    }
    childId = parentId;
  }
  
  while (1) 
  {
    float z = input_k[i];
    float *nodeWeight = weight+
    if(buffer[threadIdx.x] < z)
      buffer[threadIdx.x] = z;
  }

}

static int cunnx_SoftMaxTree_updateOutput(lua_State *L)

  THIntTensor *childParent = (THIntTensor*)luaT_getfieldcheckudata(L, 1, "childParent", "torch.IntTensor");
  THIntTensor *parentChildren = (THIntTensor*)luaT_getfieldcheckudata(L, 1, "parentChildren", "torch.IntTensor");
  
  THCudaTensor *linearOutput = luaT_getfieldcheckudata(L, 1, "_linearOutput", "torch.CudaTensor");
  THCudaTensor *logsoftOutput = luaT_getfieldcheckudata(L, 1, "_logsoftmaxOutput", "torch.CudaTensor");
  
  THCudaTensor *weight = luaT_getfieldcheckudata(L, 1, "weight", "torch.CudaTensor");
  THCudaTensor *bias = luaT_getfieldcheckudata(L, 1, "bias", "torch.CudaTensor");
