   
   luaL_argcheck(L, gradOutput->nDimension == 1, 2, "1D tensor expected");
   
-  input = THCudaTensor_newContiguous(input);
-  THCudaTensor_resize1d(output, input->size[0]);
-  THCudaTensor_zero(output);
+  THCudaTensor_resize2d(gradInput, input->size[0], input->size[1]);
   
   /* call cudakernel */
   dim3 blocks(input->size[0]); // each block is an example
   dim3 threads(SOFTMAXTREE_THREADS);
-  cunnx_SoftMaxTree_updateOutput_kernel<<<blocks,threads>>>(
+  cunnx_SoftMaxTree_updateGradInput_kernel<<<blocks,threads>>>(
     THCudaTensor_data(gradInput), THCudaTensor_data(logsoftOutput), 
     THCudaTensor_data(gradOutput), THCudaTensor_data(weight), 
     THCudaTensor_data(bias), THCudaTensor_data(target), 

  if(errcode != cudaSuccess)
    THError(cudaGetErrorString(errcode));
  
  //THCudaTensor_free(input);
  return 1;
}


  }
}



static int cunnx_SoftMaxTree_updateOutput(lua_State *L)
{ 
  THCudaTensor *input = (THCudaTensor*)luaT_checkudata(L, 2, "torch.CudaTensor");  
  THCudaTensor *target = (THCudaTensor*)luaT_checkudata(L, 3, "torch.CudaTensor");  
  int inputSize = luaT_getfieldcheckint(L, 1, "inputSize");
  int rootId = luaT_getfieldcheckint(L, 1, "rootId") - 1;
