diff --git a/abogen/is_nvidia.py b/abogen/is_nvidia.py index 2777f7b..78bd8cf 100644 --- a/abogen/is_nvidia.py +++ b/abogen/is_nvidia.py @@ -6,11 +6,15 @@ def check(): except Exception: return False + nvidia_keywords = ['nvidia', 'rtx', 'gtx', 'quadro', 'tesla', 'titan', 'mx'] for gpu in stats.gpus: - print(gpu.name) - if 'nvidia' in gpu.name.lower(): + name = gpu.name.lower() + if any(keyword in name for keyword in nvidia_keywords): return True return False if __name__ == "__main__": + stats = gpustat.new_query() + for gpu in stats.gpus: + print(gpu.name) print(check()) \ No newline at end of file