 template <typename KDType, typename VDType, typename xpu>
 inline typename std::enable_if<std::is_same<xpu, gpu>::value, size_t>::type
 SortByKeyWorkspaceSize(const size_t num_keys) {
+#ifdef SORT_WITH_THRUST
+  return 0;
+#else
   size_t sortpairs_bytes = 0;
   cub::DeviceRadixSort::SortPairs<KDType, VDType>(NULL, sortpairs_bytes,
       NULL, NULL, NULL, NULL, num_keys);
   size_t keys_bytes = num_keys*sizeof(KDType);
   size_t values_bytes = num_keys*sizeof(VDType);
   return (keys_bytes + values_bytes + sortpairs_bytes);
+#endif
 }
 
 template<typename KDType, typename VDType>

        thrust::cuda::par.on(stream),
        key_iter, key_iter + keys.size(0), value_iter, thrust::greater<KDType>());
    }
  }
  MSHADOW_CUDA_POST_KERNEL_CHECK(SortByKey);
#else
  LOG(FATAL) << "SortByKey is only supported for CUDA version >=7.0!";

  CHECK_EQ(values.CheckContiguous(), true);
#if CUDA_VERSION >= 7000
  cudaStream_t stream = mshadow::Stream<gpu>::GetStream(keys.stream_);
#ifndef SORT_WITH_THRUST
  if (workspace != NULL) {
    // Workspace given, sort using CUB
    CHECK_EQ(workspace->CheckContiguous(), true);
