Små bug fix