 
 #pragma unroll
   for (unsigned int stride = Power2SortSize / 2; stride > 0; stride /= 2) {
-    // Single warp per slice is completely synchronous
-    if (Power2SortSize > 64) {
-      __syncthreads();
-    }
-
+    
+    __syncthreads();
+    
     unsigned int pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
     bitonicSwap<Comparator, K, V>(
       keys[pos], values[pos], valid[pos],

      false, comp);
  }

  // Single warp per slice is completely synchronous
  if (Power2SortSize > 64) {
    __syncthreads();
  }
}

// Sorts (key, value) pairs (in different tensors) in-place; i.e.,

      false, comp);
  }

  __syncthreads();
  
}

template <typename Comparator, typename K,
