Commit 49a6c8c1 authored by OlivierDehaene's avatar OlivierDehaene
Browse files

fix(launcher): parse num cuda devices from CUDA_VISIBLE_DEVICES and NVIDIA_VISIBLE_DEVICES

parent 146e72c3
......@@ -455,9 +455,10 @@ fn shutdown_shards(shutdown: Arc<Mutex<bool>>, shutdown_receiver: &mpsc::Receive
}
fn num_cuda_devices() -> Option<usize> {
let devices = env::var("CUDA_VISIBLE_DEVICES")
.map_err(|_| env::var("NVIDIA_VISIBLE_DEVICES"))
.ok()?;
let devices = match env::var("CUDA_VISIBLE_DEVICES") {
Ok(devices) => devices,
Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
};
let n_devices = devices.split(',').count();
Some(n_devices)
}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment