Did a quick bugfix, hope this works.