mypy error fix

aimclub · Jan 8, 2024 · ec9867f · ec9867f
1 parent 71000c9
commit ec9867f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/stable_gnn/model_link_predict.py b/stable_gnn/model_link_predict.py
@@ -1,5 +1,5 @@
 import random
-from typing import List, Tuple
+from typing import Any, List, Tuple
 
 import torch
 from sklearn.base import BaseEstimator
@@ -38,7 +38,7 @@ def __init__(
 
     def train_test_edges(
         self, dataset: Graph
-    ) -> Tuple[List[List[int]], List[List[int]], List[List[int]], List[List[int]]]:
+    ) -> Tuple[List[Any], List[int], List[Any], List[int]]:
         """
         Split dataset to train and test and calculate negative samples