refactor: use split as splitter in example.py

IBCNServices · Aug 28, 2020 · 3fe8c4a · 3fe8c4a
1 parent f89ee53
commit 3fe8c4a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/example.py b/example.py
@@ -50,8 +50,8 @@ def create_embeddings(kg, entities, split, walkers, sg=1):
     transformer = RDF2VecTransformer(walkers=walkers, sg=sg)
     walk_embeddings = transformer.fit_transform(kg, entities)
     return (
-        walk_embeddings[: len(train_entities)],
-        walk_embeddings[len(train_entities) :],
+        walk_embeddings[:split],
+        walk_embeddings[split:],
     )