1
0

use ignite for exp2

This commit is contained in:
2025-12-02 23:07:27 +08:00
parent 43b807679f
commit 65c56e938c
15 changed files with 246 additions and 794 deletions

17
gpu_utils.py Normal file
View File

@@ -0,0 +1,17 @@
import torch
def print_gpu_availability():
"""打印PyTorch的GPU可用性"""
if torch.cuda.is_available():
print(f"GPU可用{torch.cuda.get_device_name(0)}")
else:
print("GPU不可用")
def get_gpu_device() -> torch.device:
"""获取PyTorch的GPU设备"""
if torch.cuda.is_available():
return torch.device("cuda")
else:
raise Exception("找不到CUDA")