       false, comp);
   }
 
-  // Single warp per slice is completely synchronous
-  if (Power2SortSize > 64) {
-    __syncthreads();
-  }
+  __syncthreads();
+  
 }
 
 // Sorts (key, value) pairs (in different tensors) in-place; i.e.,


#pragma unroll
  for (unsigned int stride = Power2SortSize / 2; stride > 0; stride /= 2) {
    // Single warp per slice is completely synchronous
    if (Power2SortSize > 64) {
      __syncthreads();
    }

    unsigned int pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
    bitonicSwap<Comparator, K, V>(
      keys[pos], values[pos], valid[pos],


#pragma unroll
    for (unsigned int stride = size / 2; stride > 0; stride /= 2) {
    
      __syncthreads();
      
      unsigned int pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
      bitonicSwap<Comparator, K, V>(
        keys[pos], values[pos], valid[pos],
