Use SGD for stop -token layer

2018-05-13 06:35:11 -07:00 · 2018-05-13 06:35:11 -07:00 · 074d1c97d3
parent f27281d529
commit 074d1c97d3
1 changed files with 2 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -361,7 +361,8 @@ def main(args):
                     c.r)

    optimizer = optim.Adam(model.parameters(), lr=c.lr)
-    optimizer_st = optim.Adam(model.decoder.stopnet.parameters(), lr=c.lr)
+    optimizer_st = optim.SGD(model.decoder.stopnet.parameters(), lr=0.01,
+                             momentum=0.9, nesterov=True)

    criterion = L1LossMasked()
    criterion_st = nn.BCELoss()