Fix CUDA test hang: (microsoft#6138)

- Make condition check in `CUDAAllocatorTest` to ensure CUDA device is present.
pranav-prakash · Dec 16, 2020 · 0978d2b · 0978d2b
1 parent b648bf6
commit 0978d2b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/onnxruntime/test/framework/cuda/allocator_cuda_test.cc b/onnxruntime/test/framework/cuda/allocator_cuda_test.cc
@@ -12,6 +12,10 @@ namespace onnxruntime {
 namespace test {
 TEST(AllocatorTest, CUDAAllocatorTest) {
   OrtDevice::DeviceId cuda_device_id = 0;
+
+  // ensure CUDA device is avaliable.
+  CUDA_CALL_THROW(cudaSetDevice(cuda_device_id));
+
   AllocatorCreationInfo default_memory_info(
       {[](OrtDevice::DeviceId id) { return onnxruntime::make_unique<CUDAAllocator>(id, CUDA); }, cuda_device_id});