File size: 456 Bytes
127cf1e
70a220d
 
 
2bba75d
 
0a0582d
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
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版本
# 在机器中输入命令nvidia-smi
import subprocess
try:
    result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    print('nvidia-smi output:', result.stdout)
except Exception as e:
    print(f'Error: {e}')