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
|
|
}
|
|
}
|
|
}
|