File size: 537 Bytes
127cf1e
70a220d
 
 
2bba75d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import torch
device = f'cuda:0' if torch.cuda.is_available() else 'cpu'
print('device:', device)
torch_version = torch.__version__
print('torch_version:', torch_version)

# 获取服务器cuda版本
def get_cuda_version():
    try:
        import torch
        cuda_version = torch.version.cuda
        if cuda_version is None:
            return "CUDA version not found"
        return cuda_version
    except ImportError:
        return "PyTorch is not installed"

cuda_version = get_cuda_version()
print('cuda_version:', cuda_version)