diff --git a/infer-web.py b/infer-web.py index 6443a9d..36fb09a 100644 --- a/infer-web.py +++ b/infer-web.py @@ -26,7 +26,7 @@ else: if_gpu_ok = False for i in range(ngpu): gpu_name=torch.cuda.get_device_name(i) - if("16"in gpu_name or "MX"in gpu_name):continue + if(("16"in gpu_name and "V100"not in gpu_name) or "MX"in gpu_name):continue if("10"in gpu_name or "20"in gpu_name or "30"in gpu_name or "40"in gpu_name or "A50"in gpu_name.upper() or "70"in gpu_name or "80"in gpu_name or "90"in gpu_name or "M4"in gpu_name or "T4"in gpu_name or "TITAN"in gpu_name.upper()):#A10#A100#V100#A40#P40#M40#K80 if_gpu_ok=True#至少有一张能用的N卡 gpu_infos.append("%s\t%s"%(i,gpu_name))