Fixing some bug :)