   dim3 blocks(input->size[0]); // each block is an example
   dim3 threads(SOFTMAXTREE_THREADS);
   cunnx_SoftMaxTree_accGradParameters_kernel<<<blocks,threads>>>(
-    THCudaTensor_data(gradWeight), THCudaTensor_data(Bias), 
+    THCudaTensor_data(gradWeight), THCudaTensor_data(gradBias), 
     THCudaTensor_data(input), THCudaTensor_data(linearGradOutput), 
     THCudaTensor_data(nodeUpdateCuda), THCudaTensor_data(target), 
     THCudaTensor_data(childParent), THCudaTensor_data(parentChildren), 
     input->size[1], rootId, maxFamilyPath, maxDept, scale 
   );
   
-  // copy nodeIds from device to host
-  
-  /* updates will contain parentId (key) sum of scales (value)*/
-  lua_pushinteger(L, (int)(parentId+1));
-  lua_gettable(L, -2);
-  double count = lua_tonumber(L, -1) + scale;
-  lua_pop(L, 1);
-  
-  lua_pushinteger(L, (int)(parentId+1)); /* key */
-  lua_pushnumber(L, count); /* value */
-  lua_settable(L, -3);
+  cudaError errcode = cudaGetLastError();
+  if(errcode != cudaSuccess)
+    THError(cudaGetErrorString(errcode));
+  
+  // copy updated nodeIds from device to host
+  THIntTensor_copyCuda(nodeUpdateHost, nodeUpdateCuda);
+  nodeUpdate = THIntTensor_new();
+  
+  // fill updates table
+  for (i=0; i<nodeUpdateHost->size[0]; i++)
+  {
+    THIntTensor_select(nodeUpdate, nodeUpdateHost, 0, i);
+    
+    for (j=0; j<nodeUpdateHost->size[1]; j++)
+    {
+      int nodeId = THIntTensor_get1d(nodeUpdate, j);
+      double count;
+      
+      if (nodeId == 0)
+      {
+        break;
+      }
+      
+      /* updates will contain nodeId (key) sum of scales (value)*/
+      lua_pushinteger(L, (int)(nodeId+1));
+      lua_gettable(L, -2);
+      count = lua_tonumber(L, -1) + scale;
+      lua_pop(L, 1);
+      
+      lua_pushinteger(L, (int)(nodeId+1)); /* key */
+      lua_pushnumber(L, count); /* value */
+      lua_settable(L, -3);
+    }
+  }
     
   return 0;
 }

    /* get next Node in Tree */
    node = childParent + childId*2;
    parentId = (int)node[0] - 1;
    childIdx = (int)node[1] - 1;
    
    node = parentChildren + parentId*2;
    parentIdx = (int)node[0] - 1;

      for (int j=0; j<nChildren; j++)
      {
        // multiply accumulate weights
        nodeGradWeight[j*nInput + i] += scale*nodeGradOutput[j]*buffer[tx];
      }
    }
    
    // cadd bias
    for (int j=tx; j<nChildren; j+=i_step)
    {
      // multiply accumulate biases
      nodeGradBias[j] += scale*nodeGradOutput[j];
    }
    
    // keep track of which node gets gradients
