Fixing lost class bug