Trying to fix class issues :(