Fixing bug with the graph