apply PyTorch AMP

2023-10-15 23:04:14 +09:00 · 2023-10-15 23:04:14 +09:00 · 5279516d5a
parent af154fed8d
commit 5279516d5a
1 changed files with 1 additions and 1 deletions
--- a/finetune/facenet_transfer_learning.py
+++ b/finetune/facenet_transfer_learning.py
@ -126,7 +126,7 @@ optimizer = Adam(params=[
    {'params': model_gpu[1].parameters(), 'lr': 1e-3},
 ])

-scaler = GradScaler(init_scale=256)
+scaler = GradScaler()

 # model, optimizer = optimize(model=model, optimizer=optimizer)
 scheduler = lr_scheduler.StepLR(optimizer=optimizer, step_size=10, gamma=0.9)