Took code from Will's branch to fix Travis failure