Add call to auto_set_device

rapidsai · May 12, 2022 · 7b3dd82 · 7b3dd82
1 parent ffb9c66
commit 7b3dd82
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/java/src/main/native/src/CudaJni.cpp b/java/src/main/native/src/CudaJni.cpp
@@ -392,6 +392,7 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cuda_profilerStop(JNIEnv *env, jclass
 
 JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cuda_deviceSynchronize(JNIEnv *env, jclass clazz) {
   try {
+    cudf::jni::auto_set_device(env);
     CUDF_CUDA_TRY(cudaDeviceSynchronize());
   }
   CATCH_STD(env, );