Bug fixes

Fixes issue #55 
This shouldn't be in here...?