fix(gpu): fix multi-gpu error in division

zama-ai · Aug 1, 2024 · 20c3d2e · 20c3d2e
1 parent 5547d92
commit 20c3d2e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/backends/tfhe-cuda-backend/cuda/include/integer.h b/backends/tfhe-cuda-backend/cuda/include/integer.h
@@ -2796,7 +2796,7 @@ template <typename Torus> struct int_div_rem_memory {
   int_div_rem_memory(cudaStream_t *streams, uint32_t *gpu_indexes,
                      uint32_t gpu_count, int_radix_params params,
                      uint32_t num_blocks, bool allocate_gpu_memory) {
-    active_gpu_count = get_active_gpu_count(num_blocks, gpu_count);
+    active_gpu_count = get_active_gpu_count(2 * num_blocks, gpu_count);
 
     this->params = params;
     shift_mem_1 = new int_logical_scalar_shift_buffer<Torus>(