Add back nesterov momentum

pytorch · Apr 1, 2022 · 8455b37 · 8455b37
1 parent ff6d641
commit 8455b37
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/references/detection/train.py b/references/detection/train.py
@@ -209,7 +209,13 @@ def main(args):
 
     opt_name = args.opt.lower()
     if opt_name.startswith("sgd"):
-        optimizer = torch.optim.SGD(parameters, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+        optimizer = torch.optim.SGD(
+            parameters,
+            lr=args.lr,
+            momentum=args.momentum,
+            weight_decay=args.weight_decay,
+            nesterov="nesterov" in opt_name,
+        )
     elif opt_name == "adamw":
         optimizer = torch.optim.AdamW(parameters, lr=args.lr, weight_decay=args.weight_decay)
     else: