Fixing smalll issues