Commit 64556028 authored by Gaëtan Caillaut's avatar Gaëtan Caillaut
Browse files

Adam optimizer instead of SGD

parent 6af2af9c
......@@ -133,7 +133,8 @@ if __name__ == "__main__":
if args.gpu:
model = model.to("cuda")
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr)
optimizer = torch.optim.Adam(model.parameters())
x_eval, y_eval = build_eval_batches(sequences, voc2idx, mask_idx, args.gpu)
x_eval_test, y_eval_test = build_eval_batches(
sequences_test, voc2idx, mask_idx, args.gpu)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment