Minor fix and improvements