Minor changes to classes