fix(gpu): fix default multi-bit PBS for multi-device execution of int…

…eger ops
zama-ai · Jan 28, 2025 · 65b35f0 · 65b35f0
1 parent 8eb4d86
commit 65b35f0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/backends/tfhe-cuda-backend/cuda/src/pbs/programmable_bootstrap_multibit.cuh b/backends/tfhe-cuda-backend/cuda/src/pbs/programmable_bootstrap_multibit.cuh
@@ -494,6 +494,7 @@ __host__ void execute_compute_keybundle(
     pbs_buffer<Torus, MULTI_BIT> *buffer, uint32_t num_samples,
     uint32_t lwe_dimension, uint32_t glwe_dimension, uint32_t polynomial_size,
     uint32_t grouping_factor, uint32_t level_count, uint32_t lwe_offset) {
+  cudaSetDevice(gpu_index);
 
   auto lwe_chunk_size = buffer->lwe_chunk_size;
   uint32_t chunk_size =
@@ -507,7 +508,6 @@ __host__ void execute_compute_keybundle(
       get_buffer_size_full_sm_multibit_programmable_bootstrap_keybundle<Torus>(
           polynomial_size);
   int max_shared_memory = cuda_get_max_shared_memory(gpu_index);
-  cudaSetDevice(gpu_index);
 
   auto d_mem = buffer->d_mem_keybundle;
   auto keybundle_fft = buffer->keybundle_fft;