mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-04-23 16:02:10 +00:00
afc7ded84f
Remove compute capability lock We are only calling the `get_cuda_capability` function once, so avoiding the cost of multiple calls is not really necessary yet.
22 lines
778 B
Rust
22 lines
778 B
Rust
pub fn get_cuda_capability() -> Option<(usize, usize)> {
|
|
use pyo3::prelude::*;
|
|
|
|
let py_get_capability = |py: Python| -> PyResult<(isize, isize)> {
|
|
let torch = py.import_bound("torch.cuda")?;
|
|
let get_device_capability = torch.getattr("get_device_capability")?;
|
|
get_device_capability.call0()?.extract()
|
|
};
|
|
|
|
match pyo3::Python::with_gil(py_get_capability) {
|
|
Ok((major, minor)) if major < 0 || minor < 0 => {
|
|
tracing::warn!("Ignoring negative GPU compute capabilities: {major}.{minor}");
|
|
None
|
|
}
|
|
Ok((major, minor)) => Some((major as usize, minor as usize)),
|
|
Err(err) => {
|
|
tracing::warn!("Cannot determine GPU compute capability: {}", err);
|
|
None
|
|
}
|
|
}
|
|
}
|