       for (int j=0; j<nChildren; j++)
       {
         // multiply accumulate weights
-        atomicAdd(&nodeGradWeight[j*nInput + i], scale*nodeGradOutput[j]*buffer[tx]);
-        CudaAssert(isfinite(nodeGradWeight[j*nInput + i]))
+        float dw = scale*nodeGradOutput[j]*buffer[tx];
+        CudaAssert(isfinite(dw))
+        atomicAdd(&nodeGradWeight[j*nInput + i], dw);
       }
     }
     

    // logadd?
    float max_k = buffer[SOFTMAXTREE_THREADS];
    buffer[tx] = 0;
    for (int i=tx; i<nOutput; i+=i_step)
    {
      buffer[tx] += __expf(linearOutput[i]-max_k);
      CudaAssert(isfinite(buffer[tx]))
    }

    __syncthreads();

    // reduce
    if (tx == 0)
    {
      float logsum_k = 0;
      for (int i=0; i<nOutput; i++)
        logsum_k += buffer[i];
      
      buffer[SOFTMAXTREE_THREADS] = max_k + __logf(logsum_k);
      CudaAssert(isfinite(buffer[SOFTMAXTREE_THREADS]))
    }

    __syncthreads();

    // logsoftmax
    float logsum_k = buffer[SOFTMAXTREE_THREADS];
    for (int i=tx; i<nOutput; i+=i_step)
    {
      nodeOutput[i] = linearOutput[i] - logsum_k;
      CudaAssert(isfinite(nodeOutput[i]))


    for(int i=tx; i<nChildren; i+=i_step)
    {
      nodeGrad[i] = -expf(nodeGrad[i])*grad;
      CudaAssert(isfinite(nodeGrad[i]))
    }
    
