ai-station-code/wudingpv/taihuyuan_pv/mitunet/TestAndPredict/test/test_cuda.py

5 lines
224 B
Python

import torch
print(torch.cuda.is_available()) # 检查是否有可用的 CUDA 设备
print(torch.cuda.device_count()) # 打印可用 GPU 的数量
print(torch.cuda.get_device_name(0)) # 打印第一个 GPU 的名称