github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-10 15:23:32 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

update devicelib/nnlib to detect compute capability

Browse source
This commit is contained in:
iperov 2019-01-08 10:48:06 +04:00
parent 29c2375f5a
commit e8620919a7
5 changed files with 70 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

15
nnlib/pynvml.py
View file

@ -1699,3 +1699,18 @@ def nvmlDeviceGetTopologyCommonAncestor(device1, device2):
ret = fn(device1, device2, byref(c_level))
_nvmlCheckReturn(ret)
return c_level.value
#DeepFaceLab additions
def nvmlDeviceGetCudaComputeCapability(device):
c_major = c_int()
c_minor = c_int()
fn = _nvmlGetFunctionPointer("nvmlDeviceGetCudaComputeCapability")
# get the count
ret = fn(device, byref(c_major), byref(c_minor))
# this should only fail with insufficient size
if (ret != NVML_SUCCESS):
raise NVMLError(ret)
return c_major.value, c_minor.value