From 074d1c97d334d64b1b80dc8b7370f834c6e47f55 Mon Sep 17 00:00:00 2001 From: Eren Golge Date: Sun, 13 May 2018 06:35:11 -0700 Subject: [PATCH] Use SGD for stop -token layer --- train.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/train.py b/train.py index 7ca06ed8..791c5d4f 100644 --- 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()