diff --git a/lib/memory.js b/lib/memory.js
index d126517..75473de 100644
--- a/lib/memory.js
+++ b/lib/memory.js
@@ -89,7 +89,7 @@ module.exports.startMemoryMeasurement = function (test, options, time, graphData
             }
 
             if (graphFilename && options.generateGraph) {
-                browser.controlFlow().await(
+                browser.controlFlow().wait(
                     generateGraph.phantomRender(graphFilename, graphData, options.graphWidth, options.graphHeight)
                 ).then(function () {
                         console.log('Graph generated: ' + graphFilename);