diff --git a/drivers/media/media-entity.c b/drivers/media/media-entity.c
index 9bf96c7..85af715 100644
--- a/drivers/media/media-entity.c
+++ b/drivers/media/media-entity.c
@@ -383,7 +383,13 @@
 
 	mutex_lock(&mdev->graph_mutex);
 
-	media_entity_graph_walk_start(graph, entity);
+	ret = media_entity_graph_walk_init(&pipe->graph, mdev);
+	if (ret) {
+		mutex_unlock(&mdev->graph_mutex);
+		return ret;
+	}
+
+	media_entity_graph_walk_start(&pipe->graph, entity);
 
 	while ((entity = media_entity_graph_walk_next(graph))) {
 		DECLARE_BITMAP(active, MEDIA_ENTITY_MAX_PADS);
@@ -481,6 +487,8 @@
 			break;
 	}
 
+	media_entity_graph_walk_cleanup(graph);
+
 	mutex_unlock(&mdev->graph_mutex);
 
 	return ret;
@@ -502,6 +510,8 @@
 			entity->pipe = NULL;
 	}
 
+	media_entity_graph_walk_cleanup(graph);
+
 	mutex_unlock(&mdev->graph_mutex);
 }
 EXPORT_SYMBOL_GPL(media_entity_pipeline_stop);