     for (int j=tx; j<nChildren; j+=i_step)
     {
       // multiply accumulate biases
-      atomicAdd(&nodeGradBias[j], scale*nodeGradOutput[j]);
-      CudaAssert(isfinite(nodeGradBias[j]))
+      float db = scale*nodeGradOutput[j];
+      CudaAssert(isfinite(db))
+      atomicAdd(&nodeGradBias[j], db);
     }
     
     // keep track of which node gets gradients

      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]))
      }
    }
    

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

    // reduce
    for (unsigned int stride = blockDim.x >> 1; stride > 0; stride >>= 1)
    {
      __syncthreads();
      if (tx < stride)
        buffer[tx] += buffer[tx+stride];
    }
    if (tx == 0)
    {
      float m = max_k + logf(buffer[0]);
      CudaAssert(isfinite(m))
      buffer[SOFTMAXTREE_THREADS] = m;
    }

    __syncthreads();

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