songsenand
/
SUimeModelTraner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(trainer): 检测GPU可用性并提示用户退回CPU训练

This commit is contained in:
songsenand 2026-04-05 00:32:12 +08:00
parent 4823afcd1f
commit b3055656d1
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/model/trainer.py
View File

@ -812,7 +812,8 @@ def train(
prefetch_factor=64, # 每个worker预取64个batch prefetch_factor=64, # 每个worker预取64个batch
persistent_workers=True, # 保持worker进程存活 persistent_workers=True, # 保持worker进程存活
) )
if torch.cuda.is_available():
console.print("[red]GPU不可用，退回使用CPU[/red]")
console.print("[green]✓ 数据加载器创建完成[/green]") console.print("[green]✓ 数据加载器创建完成[/green]")
console.print(f" 训练批次大小: {batch_size}") console.print(f" 训练批次大小: {batch_size}")
console.print(f" 评估批次大小: {batch_size}") console.print(f" 评估批次大小: {batch_size}")