summaryrefslogtreecommitdiff
path: root/.tmp_gpu_check.py
blob: 706491cb4b3da104e1859b87e4cc9229146a8c66 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
import os
print("ENV CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES"))
try:
    import torch
    print("torch version:", torch.__version__)
    print("torch.version.cuda:", getattr(torch.version, "cuda", None))
    print("torch.cuda.is_available():", torch.cuda.is_available())
    print("torch.cuda.device_count():", torch.cuda.device_count())
    if torch.cuda.is_available():
        for i in range(torch.cuda.device_count()):
            print(f"[{i}]", torch.cuda.get_device_name(i))
except Exception as e:
    print("torch error:", repr(e))