Fixed bug in last commit, seems to work now.