Fix for issue 5
Cleanup