FIX small bugs