Minor fixes to class