fix #5

Just a quick fix. All tests pass, but I am not really familiar with the code so I am not sure if it breaks anything.