Added nesterov momentum

ancestor-mithril · Nov 7, 2024 · 4335623 · 4335623
1 parent 7076751
commit 4335623
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/utils/optimizer.py b/utils/optimizer.py
@@ -11,6 +11,7 @@ def init_optimizer(args, parameters) -> Optimizer:
         "lr": args.lr,
         "momentum": momentum,
         "weight_decay": weight_decay,
+        "nesterov": True,
     }
     if torch.torch_version.TorchVersion(torch.__version__) >= "2.3.0":
         if args.device != "cpu":