Update

kumo-ai · Sep 30, 2024 · 9bc4dea · 9bc4dea
1 parent 08bf337
commit 9bc4dea
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/examples/gnn_link.py b/examples/gnn_link.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import argparse
 import copy
 import json
@@ -245,7 +247,7 @@ def test(src_loader: NeighborLoader, dst_loader: NeighborLoader) -> np.ndarray:
             best_val_metric = val_metrics[tune_metric]
             state_dict = copy.deepcopy(model.state_dict())
 
-model.load_state_dict(state_dict)
+model.load_state_dict(state_dict)  # type: ignore
 val_pred = test(*eval_loaders_dict["val"])
 val_metrics = task.evaluate(val_pred, task.get_table("val"))
 print(f"Best Val metrics: {val_metrics}")