print

tfhe/src/core_crypto/gpu/mod.rs

@@ -53,6 +53,7 @@ impl CudaStreams {
             ptr_array.push(unsafe { cuda_create_stream(i.get()) });
             gpu_indexes.push(i);
         }
+        println!("GPU indexes for streams: {:?}", gpu_indexes);
         Self {
             ptr: ptr_array,
             gpu_indexes,

tfhe/src/high_level_api/tests/gpu_selection.rs

@@ -152,15 +152,16 @@ fn test_specific_gpu_selection() {
     assert_eq!(a.gpu_indexes(), &[]);
     assert_eq!(b.gpu_indexes(), &[]);
 
+    let gpus = gpus_to_be_used.clone();
+    let used_gpus = gpus.gpu_indexes();
     let cuda_key = compressed_server_keys.decompress_to_specific_gpu(gpus_to_be_used);
 
-    let first_gpu = GpuIndex::new(0);
-
     set_server_key(cuda_key);
+    println!("Here 0");
     let c = &a + &b;
     let decrypted: u32 = c.decrypt(&keys);
     assert_eq!(c.current_device(), Device::CudaGpu);
-    assert_eq!(c.gpu_indexes(), &[first_gpu]);
+    assert_eq!(c.gpu_indexes().first().unwrap(), used_gpus.first().unwrap());
     assert_eq!(decrypted, clear_a.wrapping_add(clear_b));
 
     // Check explicit move, but first make sure input are on Cpu still
@@ -174,12 +175,13 @@ fn test_specific_gpu_selection() {
 
     assert_eq!(a.current_device(), Device::CudaGpu);
     assert_eq!(b.current_device(), Device::CudaGpu);
-    assert_eq!(a.gpu_indexes(), &[first_gpu]);
-    assert_eq!(b.gpu_indexes(), &[first_gpu]);
+    assert_eq!(a.gpu_indexes().first().unwrap(), used_gpus.first().unwrap());
+    assert_eq!(b.gpu_indexes().first().unwrap(), used_gpus.first().unwrap());
 
+    println!("Here 1");
     let c = &a + &b;
     let decrypted: u32 = c.decrypt(&keys);
     assert_eq!(c.current_device(), Device::CudaGpu);
-    assert_eq!(c.gpu_indexes(), &[first_gpu]);
+    assert_eq!(c.gpu_indexes().first().unwrap(), used_gpus.first().unwrap());
     assert_eq!(decrypted, clear_a.wrapping_add(clear_b));
 }