                      nBlocksPerSeg);
             DEBUG_SYNC;
             segmentedSum(NULL, NULL, (V*)temp_storage, out,
-                         Linear(nBlocksPerSeg, 0), nBlocksPerSeg, nSegments, devNo, stream);
+                         Linear(nBlocksPerSeg, 0), nBlocksPerSeg, nSegments, nThreadsToFillDevice, stream);
         }
     }
 }

    }
    else {
        /* General case */
        int nThreadsToFillDevice = getNumThreadsToFillDevice(devNo);
        int nThreads = divru(segLen, 128) * nSegments;
        int nLoops = divru(nThreads, nThreadsToFillDevice);
        if (2 <= nLoops) {

void segmentedSum(void *temp_storage, sq::SizeType *temp_storage_bytes,
                  InputIterator in, OutputIterator out,
                  OffsetIterator segOffset, int segLen,
                  sq::SizeType nSegments, int nThreadsToFillDevice, cudaStream_t stream) {
    *temp_storage_bytes = 0;
    if (segLen <= WARP_SIZE) {
        dim3 blockDim(128);
