apply PyTorch AMP
continuous-integration/drone/push Build was killed Details

This commit is contained in:
yayoimizuha 2023-10-15 22:50:55 +09:00
parent 4d9595fe5e
commit 1d943d029b
1 changed files with 1 additions and 1 deletions

View File

@ -126,7 +126,7 @@ optimizer = Adam(params=[
{'params': model_gpu[1].parameters(), 'lr': 1e-3}, {'params': model_gpu[1].parameters(), 'lr': 1e-3},
]) ])
scaler = GradScaler(init_scale=4096) scaler = GradScaler(init_scale=256)
# model, optimizer = optimize(model=model, optimizer=optimizer) # model, optimizer = optimize(model=model, optimizer=optimizer)
scheduler = lr_scheduler.StepLR(optimizer=optimizer, step_size=10, gamma=0.9) scheduler = lr_scheduler.StepLR(optimizer=optimizer, step_size=10, gamma=0.9)