Bugs correction (#4)

Fixing issue #4