Fixed small bugs after merging