add pretrain-mode to gats

2026-07-02 10:31:00 +08:00 · 2020-11-25 17:46:34 +08:00
parent c14a99a735
commit cd7c81cfd0
5 changed files with 27 additions and 4 deletions
--- a/examples/workflow_by_code_gats.py
+++ b/examples/workflow_by_code_gats.py
@@ -70,9 +70,10 @@ if __name__ == "__main__":
                "lr": 1e-3,
                "early_stop": 20,
                "batch_size": 800,
-                "metric": "IC",
+                "metric": "loss",
                "loss": "mse",
-                "base_model": "GRU",
+                "base_model": "LSTM",
+                "with_pretrain": True,
                "seed": 0,
                "GPU": 0,
            },