Merge pull request #3 from martinnovaak/adamw

martinnovaak · web-flow · commit 586058f619d8 · 2024-05-05T21:10:52.000+02:00
Use AdamW instead of Adam
diff --git a/trainer/main.py b/trainer/main.py
@@ -34,7 +34,7 @@ def load_config(config_path="config.json"):
 
     dataloader = BatchLoader(data_parser_path, paths, batch_size, scale, wdl)
 
-    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-2)
 
     start_time = time()
 
