Use AdamW instead of Adam

martinnovaak · martinnovaak · commit 67123a16779d · 2024-05-05T21:10:29.000+02:00
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()
 
