Skip to content

Commit 2ac8d55

Browse files
authored
Update infer-web.py
1 parent dc0c875 commit 2ac8d55

File tree

1 file changed

+3
-5
lines changed

1 file changed

+3
-5
lines changed

infer-web.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,6 @@
2929
if_gpu_ok = False
3030
for i in range(ngpu):
3131
gpu_name = torch.cuda.get_device_name(i)
32-
if ("16" in gpu_name and "V100" not in gpu_name) or "MX" in gpu_name:
33-
continue
3432
if (
3533
"10" in gpu_name
3634
or "20" in gpu_name
@@ -44,10 +42,10 @@
4442
or "70" in gpu_name
4543
or "80" in gpu_name
4644
or "90" in gpu_name
47-
or "M4" in gpu_name
48-
or "T4" in gpu_name
45+
or "M4" in gpu_name.upper()
46+
or "T4" in gpu_name.upper()
4947
or "TITAN" in gpu_name.upper()
50-
): # A10#A100#V100#A40#P40#M40#K80
48+
): # A10#A100#V100#A40#P40#M40#K80#A4500
5149
if_gpu_ok = True # 至少有一张能用的N卡
5250
gpu_infos.append("%s\t%s" % (i, gpu_name))
5351
gpu_info = (

0 commit comments

Comments
 (0)