set betas and weight decay for optimizers

according to suggestions in #118 (comment)

ghstack-source-id: d03375a1eedd660d116b25609db906eaf9cc2c50
Pull Request resolved: #123

train.py

@@ -65,9 +65,10 @@ def build_optimizer(model, job_config: JobConfig):
     name = job_config.optimizer.name
     lr = job_config.optimizer.lr
     if name == "Adam":
-        optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+        # TODO: make the optimizer options configurable by toml/cmd args
+        optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1)
     elif name == "AdamW":
-        optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+        optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1)
     else:
         raise NotImplementedError(f"optimizer {name} not added")