diff --git a/py/gp.py b/py/gp.py
index a98c6b4..4dea0e7 100644
--- a/py/gp.py
+++ b/py/gp.py
@@ -1,4 +1,6 @@
 import numpy as np
+
+from node2vec import Node2Vec
 from sklearn.gaussian_process.kernels import Kernel, Hyperparameter
 from sklearn.gaussian_process.kernels import GenericKernelMixin
 from sklearn.gaussian_process import GaussianProcessRegressor
@@ -12,7 +14,7 @@ class BookKernel(GenericKernelMixin, Kernel):
 
         self.G = G
         self.node2vec = Node2Vec(self.G, dimensions=32, walk_length=16, num_walks=256, workers=8)
-        self.model = node2vec.fit(window=10, min_count=1, batch_words=4)
+        self.model = self.node2vec.fit(window=10, min_count=1, batch_words=4)
         self.wv = self.model.wv
 
     def _f(self, s1, s2):