diff --git a/node2vec/node2vec.py b/node2vec/node2vec.py
index e2d611e..4dc7f03 100644
--- a/node2vec/node2vec.py
+++ b/node2vec/node2vec.py
@@ -115,9 +115,6 @@ def _precompute_probabilities(self):
                 d_graph[current_node][self.PROBABILITIES_KEY][
                     source] = unnormalized_weights / unnormalized_weights.sum()
 
-                # Save neighbors
-                d_graph[current_node][self.NEIGHBORS_KEY] = d_neighbors
-
             # Calculate first_travel weights for source
             first_travel_weights = []
 
@@ -127,6 +124,9 @@ def _precompute_probabilities(self):
             first_travel_weights = np.array(first_travel_weights)
             d_graph[source][self.FIRST_TRAVEL_KEY] = first_travel_weights / first_travel_weights.sum()
 
+            # Save neighbors
+            d_graph[source][self.NEIGHBORS_KEY] = list(self.graph.neighbors(source))
+
     def _generate_walks(self) -> list:
         """
         Generates the random walks which will be used as the skip-gram input.